ROCm/jax

Pull Requests Commits

apaszke committed 5 years ago

4773f329

Add support for staging out the primal computation

apaszke committed 5 years ago

1214b395

Finish tracers_to_jaxpr

apaszke committed 5 years ago

257d8175

use cleaner process_primitive implementation

mattjj committed 5 years ago

c7829385

add more reasonable linearize sketch

mattjj committed 5 years ago

3c5ce13b

shady linearize experiment

mattjj committed 5 years ago

b4436f65

Merge pull request #2993 from google/cond

froystig committed 5 years ago

Verified 450b7abb

silence a pytype error in jet tracing

froystig committed 5 years ago

b19166c4

improve two-operand cond detection

froystig committed 5 years ago

28da2bc4

Add support for sorting complex values, defaulting to a NumPy-style l… (#3096)

hawkinsp committed 5 years ago

Verified 77703b89

Allow non-classes as second argument of `issubclass`. (#3098)

hawkinsp committed 5 years ago

Verified 85d7b308

stash order on jet master trace, fixes #3079 (#3097)

mattjj committed 5 years ago

Verified dceb5787

avoid committing to argument types so long as cond is overloaded

froystig committed 5 years ago

9d8ecc5a

Adds additional epsilon to adam for numerical stability. (#3091)

thomaskeck committed 5 years ago

Verified 007cdf2f

fix astype() test (#3072)

jakevdp committed 5 years ago

Verified 83c7394b

sharded_jit cleanup (#3075)

skye committed 5 years ago

Verified 7c687b24

changelog: single-operand cond

froystig committed 5 years ago

73a55e04

update reference jaxpr using cond in outfeed rewriting test

froystig committed 5 years ago

930be588

handle single-operand cond in token-threading rewriter

froystig committed 5 years ago

cb5b1d10

update uses of cond in host callback test

froystig committed 5 years ago

9f6774a9

Make lax.sort support tuple arguments using a variadic sort. (#3085)

hawkinsp committed 5 years ago

Verified 4ce2aa25

comment on joining constants for conditional branch jaxprs

froystig committed 5 years ago

6e3bfc33

comment on joining staged jaxprs in partial evaluation of conditionals

froystig committed 5 years ago

2950eb11

style changes in lax.cond

froystig committed 5 years ago

34efdd98

remove deprecation warning on five-argument cond

froystig committed 5 years ago

e8f12b6b

test and fix cond when branch-staging is off

froystig committed 5 years ago

76612c62

update jaxpr doc and tests with single-operand cond

froystig committed 5 years ago

de03c99b

move cond operand to final argument position

froystig committed 5 years ago

f90bd4f5

update reference jaxpr in cond-related jaxpr test

froystig committed 5 years ago

cc6cea20

cache jaxprs formed by cond branch staging

froystig committed 5 years ago

df622798

Older