ROCm/jax

Pull Requests Commits

revive the tracer leak checker

mattjj committed 5 years ago

9b12763b

fix #3165 if round half up behavior is desired (#3166)

jenkspt committed 5 years ago

Verified de800d20

Update Common_Gotchas_in_JAX.ipynb (#3189)

Sohl-Dickstein committed 5 years ago

Verified 190f88de

Merge pull request #2734 from google/tycheck

froystig committed 5 years ago

Verified 96c20f32

work around mypy

froystig committed 5 years ago

c293a102

except-and-raise during jaxpr checking, adding jaxpr as context, and simplify type environment

froystig committed 5 years ago

69d7bcf7

fix unitvar comparisons and move to class attributes

froystig committed 5 years ago

8e61ce8d

Updates onp -> np in random, loops, jet and in the tests of stax and optix (#3182)

joaogui1 committed 5 years ago

Verified 77e4d8b3

Fix pxla.shard_args bug (#3170)

skye committed 5 years ago

Verified a3e0cd12

axis_index abstract eval rule

froystig committed 5 years ago

5d125539

return tuple in psum abstract eval rule

froystig committed 5 years ago

1a916626

use jax.numpy in jaxpr typecheck tests

froystig committed 5 years ago

1d780816

extend type transfer to all primitives, including call and map primitives

froystig committed 5 years ago

7ff389bd

raise type errors consistently in jaxpr checker

froystig committed 5 years ago

e2cc5689

fix import in core_test

froystig committed 5 years ago

6475f60c

avoid attempt to read literals from the typechecking environment

froystig committed 5 years ago

1e556033

tidy jaxpr typechecking error test

froystig committed 5 years ago

060bd8a4

add jaxpr context to typechecker error message

froystig committed 5 years ago

0f109d9f

have UnitVar subclass Var (caught by mypy)

froystig committed 5 years ago

3705252b

update check_jaxpr doc

froystig committed 5 years ago

42e7e20e

check aval compatibility, not strict equality, when typechecking jaxpr equations

froystig committed 5 years ago

cc34ed26

check that variables are typed equally throughout a jaxpr

froystig committed 5 years ago

0c2c5584

factor out jaxpr-check context and variable environment

froystig committed 5 years ago

8e70769c

factor out jaxpr equation checks

froystig committed 5 years ago

1205f7a0

raise TypeError for jaxpr typechecking errors

froystig committed 5 years ago

94b1f631

typecheck jaxpr equations

froystig committed 5 years ago

82a9af51

Update jax version to 0.1.68 (#3181)

skye committed 5 years ago

Verified 4cbd14ca

Future-proof view test against signaling NaNs (#3178)

jakevdp committed 5 years ago

Verified bb2127ce

Added `associative_scan`. (#2170)

dpiponi committed 5 years ago

Verified c459280e

revise vmap in_axes/out_axes leaf type error msg (#3179)

mattjj committed 5 years ago

Verified 5c1de283

Older