ROCm/jax

Pull Requests Commits

[ROCm]: Add Jax-Triton support for ROCm

Rahul Batra committed 2 years ago

4dd54a8a

Disable scipy_spatial_test::RotationInv

Rahul Batra committed 2 years ago

a99f96f9

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

023b1fd3

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

e09fb404

[ROCm]: Updates defaults in build script

Rahul Batra committed 2 years ago

c38d0eff

[ROCm]: Updates for container and build script

Rahul Batra committed 2 years ago

1e3b0958

Merge pull request #16305 from skye:version

a-googler committed 2 years ago

4c02f2c7

Remove unnecessary assignment to `num_tokens`

a-googler committed 2 years ago

9dccaad9

Update WORKSPACE and setup.py for jax + jaxlib 0.4.12 release

skye committed 2 years ago

8223e452

Delete OpShardingSharding export since it has been 3 months since it was deprecated. Also remove deprecation warnings for MeshPspecSharding.

yashk2810 committed 2 years ago

14451492

Make `result_shardings`'s documentation more accurate

a-googler committed 2 years ago

03b06b6a

[XLA:Python] Fix crash when accessing locals of JAX-generated Python tracebacks.

hawkinsp committed 2 years ago

6374b73c

[XLA:Python] Fix incorrect code source information under Python 3.11.

hawkinsp committed 2 years ago

a47aca82

[PJRT PLUGIN] Pass in a blocking key value get and a key value put function pointer instead of a DistributedRuntimeClient pointer when creating a GPU client.

Jieying Luo committed 2 years ago

e598bc8e

Merge pull request #16072 from axch:stacked-piles

a-googler committed 2 years ago

5c5baccb

Merge pull request #16276 from skye:maxfail

a-googler committed 2 years ago

237625e0

axch committed 2 years ago

00554cf6

Merge pull request #16307 from jakevdp:float8-test-util

a-googler committed 2 years ago

f08563c8

Merge pull request #16246 from chrisflesher:scipy-rotation-v3

a-googler committed 2 years ago

8d27f206

Merge pull request #16271 from jakevdp:abstract-array-deprecation

a-googler committed 2 years ago

6518e4e3

Merge pull request #16309 from jakevdp:fix-typo

a-googler committed 2 years ago

09dce74a

Make the device check works when the number of devices is larger than one.

a-googler committed 2 years ago

d9797a48

Added scipy.spatial.transform Rotation and Slerp classes

chrisflesher committed 2 years ago

5be17ed9

Typo: remove stray tick in jax.numpy docs

jakevdp committed 2 years ago

3fc70d3d

Fix test_util for new float8 type

jakevdp committed 2 years ago

8d165193

Mark jax.abstract_arrays as deprecated

jakevdp committed 2 years ago

47ae5bdd

Fix the bug when lax.gather uses dynamic_gather.

maxwillzq committed 2 years ago

0ec9f3c2

Merge pull request #16299 from mattjj:checkify-ian

a-googler committed 2 years ago

d1c4c256

fix checkify + custom_vjp after symbolic zeros change

mattjj committed 2 years ago

01fa7e07

Fix `get_tensorstore_spec` for GCS paths if ocdbt is enabled

yashk2810 committed 2 years ago

b44f8b4b

Older