ROCm/jax

Pull Requests Commits

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

Merge pull request #18928 from jakevdp:jnp-linalg

a-googler committed 2 years ago

1851447b

array api: add jnp.linalg.cross & jnp.linalg.outer

jakevdp committed 2 years ago

fe2ad892

Merge pull request #18943 from jakevdp:xla-bridge

a-googler committed 2 years ago

2fa90e1d

Merge pull request #18888 from superbobry:pp-improvement

a-googler committed 2 years ago

616f4d29

Add missing __future__ import

jakevdp committed 2 years ago

796449e3

Increase test timeout for Cloud TPU nightly CI.

Jieying Luo committed 2 years ago

8128b54e

`mesh_utils.create_hybrid_device_mesh`: make sorting granules by key user configurable.

a-googler committed 2 years ago

94d58b72

The pretty printer now de-duplicates identical jaxprs

superbobry committed 2 years ago

840abfb7

[export] Add support for serialization and deserialization of Exported

gnecula committed 2 years ago

b077483b

Newer Older