ROCm/jax

Pull Requests Commits

very sketchy vmap collectives

mattjj committed 5 years ago

e69305ef

added basic pmap handling to gmap!

mattjj committed 5 years ago

1f1122cf

check gmap jaxpr by calling impl directly

mattjj committed 5 years ago

72ad64ec

apaszke committed 5 years ago

bca8a347

Initial version of gmap

apaszke committed 5 years ago

7f36d05f

omnistaging wip

mattjj committed 5 years ago

df975574

[jax2tf] Add special case for translation of lax.gather to tf.gather. (#3486)

gnecula committed 5 years ago

Verified 496cde6e

lax.sort: stable by default

froystig committed 5 years ago

ccb640af

detailed assertion in profiler test

froystig committed 5 years ago

6eea0ce1

limit jaxpr context in typechecker error messages

froystig committed 5 years ago

c7e97b79

introduce custom typecheck rules, implement them for cond and scan

froystig committed 5 years ago

c54acbf9

raise a custom error in jaxpr checker

froystig committed 5 years ago

6b3b42d9

Correction a typo of the period of the PRNG. (#3578)

frederikwilde committed 5 years ago

Verified 42cbe49c

Add a heap profiler API and document it. (#3576)

hawkinsp committed 5 years ago

Verified 5116fd47

Clarify docs on jax.lax.cond. (#3569)

malcolmreynolds committed 5 years ago

Verified 63ff6cb8

ensure lax.reduce monoid test uses original numpy (#3573)

mattjj committed 5 years ago

Verified 11caa21e

Added missing is_stable argument to lax.sort (#3553)

nova77 committed 5 years ago

Verified 99a43f20

Update README for jaxlib 0.1.50 release. (#3574)

hawkinsp committed 5 years ago

Verified b576417e

Fix test failures on GPU. (#3572)

hawkinsp committed 5 years ago

Verified 66cea027

Update docker script for CUDA 11. (#3571)

hawkinsp committed 5 years ago

Verified 8f86b139

fix error when doing forward-mode of odeint (#3566)

mattjj committed 5 years ago

Verified 26c6c3a4

Fix logsumexp test (#3563)

jakevdp committed 5 years ago

Verified 4b1bb189

removed stale faq entries (#3565)

mattjj committed 5 years ago

Verified 062ce297

update version and changelog for pypi (#3564)

mattjj committed 5 years ago

Verified c2501d1b

allow closures for odeint dynamics functions (#3562)

mattjj committed 5 years ago

Verified db80ca5d

Update XLA in preparation for a new jaxlib release (0.1.50). (#3560)

hawkinsp committed 5 years ago

Verified 2a6fc316

Make CUDA wheels manylinux2010 compliant, add CUDA 11, drop CUDA 9.2 (#3555)

hawkinsp committed 5 years ago

Verified a141cc6e

Fix lazy broadcast issue (#3536)

j-towns committed 5 years ago

Verified c9670d50

add remat to top-level docs (#3554)

mattjj committed 5 years ago

Verified 93054fa1

Fix eigh JVP to ensure that both the primal and tangents of the eigen… (#3550)

hawkinsp committed 5 years ago

Verified 32e419d1

Older