ROCm/jax

Pull Requests Commits

[ROCm]: Lower sparse(some) ops correctly for ROCm

Rahul Batra committed 2 years ago

d7b25908

Merge pull request #18581 from 8bitmp3:jax-docs-thinking-in-jax

a-googler committed 2 years ago

29ed3cd4

Reverts 3c07c10a9a55f9a32390dd10cf3f420bdf3f1ed8

yashk2810 committed 2 years ago

b52bcc16

Upgrade How to think in JAX

8bitmp3 committed 2 years ago

a4003bf0

Fix the pjit flakey test. The test was weirdly written in the first place. The current suspicion is that `x.copy()` inside jit made it flakey. Also delete some duplicate tests from the time of migrating to jax.Array

yashk2810 committed 2 years ago

e8888065

Start jax and jaxlib 0.4.22 release

yashk2810 committed 2 years ago

2e633522

Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md.

yashk2810 committed 2 years ago

3c07c10a

Merge pull request #18869 from 8bitmp3:small-fix-jax-pytrees

a-googler committed 2 years ago

0281596e

Merge pull request #18582 from 8bitmp3:jax-docs-debugging-101

a-googler committed 2 years ago

d269c581

Merge pull request #18961 from mattjj:issue18955

a-googler committed 2 years ago

4459991d

Merge pull request #18949 from froystig:seed-offset

a-googler committed 2 years ago

196c97fa

Merge pull request #18951 from mattjj:shmap-varargs-error

a-googler committed 2 years ago

5104c6b0

[shard-map] register cumsum et al with generic rules

mattjj committed 2 years ago

4ba6bd51

[shard-map] fix varargs error message bug

mattjj committed 2 years ago

2bff2f40

Update XLA dependency to use revision

a-googler committed 2 years ago

77c08b4a

Update XLA dependency to use revision

a-googler committed 2 years ago

fff7300b

Merge pull request #18933 from superbobry:pyupgrade

a-googler committed 2 years ago

c3dba368

Upgrade remaining sources to Python 3.9

superbobry committed 2 years ago

f936613b

Import new version of Triton

a-googler committed 2 years ago

4991aac3

Update XLA dependency to use revision

a-googler committed 2 years ago

4183f29c

introduce a config flag to control a random seed offset

froystig committed 2 years ago

67179073

Update working-with-pytrees.md

8bitmp3 committed 2 years ago

7bb4b0fa

Remove the old cache-key generation code.

a-googler committed 2 years ago

32c99f62

Upgrade JAX debugging doc

8bitmp3 committed 2 years ago

14407b9a

Merge pull request #18946 from jakevdp:dynamic-shape-error

a-googler committed 2 years ago

05df8750

Add jax[cuda12] install variation for using cuda plugin.

Jieying Luo committed 2 years ago

4fe9e596

Add cuda to _nonexperimental_plugins. It passed existing GPU tests, and have set up continuous testing job.

Jieying Luo committed 2 years ago

7305b64f

increase sample sizes on noisy random tests

froystig committed 2 years ago

8dcc079d

Improve shape validation when jax_dynamic_shapes=True

jakevdp committed 2 years ago

a1ee8c17

Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.

yashk2810 committed 2 years ago

f210b0f9

Older