ROCm/jax

Pull Requests Commits

[ROCM] Fix export test for rocm

Ruturaj4 committed 1 year ago

651cc1cb

Merge pull request #21949 from hawkinsp:winwheel

a-googler committed 2 years ago

f4158ace

Add --allow-downgrade to Windows wheel builds.

hawkinsp committed 2 years ago

b0b02682

Merge pull request #21945 from hawkinsp:release

a-googler committed 2 years ago

5d35c990

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

d7bc6b4d

Prepare for 0.4.30 release.

hawkinsp committed 2 years ago

3f4f79c8

Pallas GPU no longer falls back to lax.pow for integer powers

superbobry committed 2 years ago

dfcfb360

Removed unnecessary skip in pallas_test.py::SoftmaxTest

superbobry committed 2 years ago

5bfd6afa

[jax] Enable api_test with XLA:CPU thunks

ezhulenev committed 2 years ago

3fd93268

[XLA:Mosaic] Add internal scratch VMEM

bythew3i committed 2 years ago

ed4958cb

[Pallas/TPU] Add API for megacore partitioning of pipelines

sharadmv committed 2 years ago

701c63e1

[Pallas] Add checkify support for pallas_call in interpret mode.

justinjfu committed 2 years ago

fb68f344

[Pallas] Add initial DMA interpret mode rules. Currently this only supports LOGICAL device ids with one sharding axis.

justinjfu committed 2 years ago

1d77720e

[Pallas] Fix warning for passing None into key in pallas PRNG.

justinjfu committed 2 years ago

9556a29c

[Pallas] Make num_programs return an int if the grid is not dynamic

sharadmv committed 2 years ago

9499de43

Fix build dependency in CUDA custom call example

dfm committed 2 years ago

1de2756c

Don't use deprecated `backend=cpu` in layout test

yashk2810 committed 2 years ago

2daa0af4

Merge pull request #21918 from google:dependabot/pip/scipy-1.13.1

a-googler committed 2 years ago

3044da39

Merge pull request #21645 from andportnoy:aportnoy/cuda-custom-call

a-googler committed 2 years ago

f819c344

Fix a bug in `device_put` lowering introduced by https://github.com/google/jax/pull/21754

junwhanahn committed 2 years ago

81f63ed1

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

68387832

Deprecate `jax.xla_computation`. Use JAX AOT APIs to get the equivalent of `jax.xla_computation` functionality.

yashk2810 committed 2 years ago

e6f26ff2

Bump scipy from 1.13.0 to 1.13.1

dependabot[bot] committed 2 years ago

Verified b26af709

Adjust test tolerances for Eigh and Svd to make tests pass on all platforms.

Google-ML-Automation committed 2 years ago

9e1f1388

Merge pull request #21922 from hawkinsp:deps

a-googler committed 2 years ago

b608c69b

Use NumPy 2.0.0 and SciPy 1.13.1 in builds.

hawkinsp committed 2 years ago

160e09e2

Add CUDA custom call example as a JAX test

andportnoy committed 2 years ago

ec5c4f5a

Merge pull request #21912 from superbobry:pallas

a-googler committed 2 years ago

c9b23f0f

Add `lowering_platforms` to `traced.lower()` to allow lowering to different backends and multi-backend lowering too. In other words, enable cross-lowering!

yashk2810 committed 2 years ago

6ba16e03

Merge pull request #21905 from rajasekharporeddy:doc_typos

a-googler committed 2 years ago

be1f4ba3

Older