ROCm/jax

Pull Requests Commits

Update rules_ml_toolchain to latest head

alekstheod committed 3 days ago

b1f1db2c

[Mosaic GPU] Correctly propagate `approx=True` to `FragmentedArray.ex2` in `FragmentedArray.exp(approx=True)`.

bchetioui committed 4 days ago

2fd83c70

[Pallas/Mosaic GPU] Add `__eq__` and `__hash__` implementations for `AbstractRefUnion`.

bchetioui committed 4 days ago

113cb4d2

[trace_to_jaxpr] Refactor mosaic/core.py for trace_to_jaxpr.

gnecula committed 4 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 4 days ago

c24ca37d

Drop support for NumPy 2.0 and SciPy 1.14 for the jax 0.11.0 release.

hawkinsp committed 4 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 4 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 4 days ago

59eea209

Remove stale version tests now the minimum jaxlib is v0.10.2.

hawkinsp committed 4 days ago

86a46a6e

[Mosaic:GPU] Remove NVSHMEM usage from Mosaic XLA custom kernel and tests.

PatriosTheGreat committed 4 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 4 days ago

2e4d7f65

[JAX] Update build.py on how to update the requirements files.

hawkinsp committed 4 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 4 days ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 4 days ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 4 days ago

31c559c7

Drop support for Python 3.11 for the next JAX release (0.11.0 on July 15).

hawkinsp committed 4 days ago

5ab345f1

Try to resolve test failure on ROCM.

akuegel committed 4 days ago

305df9b6

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in Pallas

gnecula committed 4 days ago

dc575083

Reverts dced42f5da959fad3268cbeedc024e0ffd8951a5

bchetioui committed 5 days ago

452abf87

Update XLA dependency to use revision http://github.com/openxla/xla/commit/c1fcf2507528eec72374aeb2eddd85d028939cc2

Google-ML-Automation committed 5 days ago

2e34221d

Make async collectives return the output of bind directly instead of wrapping it in `Future` object (which is now deleted. We'll add it back once we need an eager representation of AbstractFuture).

yashk2810 committed 5 days ago

05127a4c

Remove stale version guards after JAX v0.10.2 release.

hawkinsp committed 5 days ago

55d755f4

[JAX] Add `py_pjrt_host_callback` extension to jaxlib

Google-ML-Automation committed 5 days ago

cf6655fd

Merge pull request #38550 from hawkinsp:postrelease

Google-ML-Automation committed 5 days ago

6d4aff5f

Updates after JAX release 0.10.2

hawkinsp committed 5 days ago

8e784c54

Merge pull request #38538 from hawkinsp:continuoustest

Google-ML-Automation committed 5 days ago

c63bb6c1

Merge pull request #38405 from skye-ant:cache_host_callbacks

Google-ML-Automation committed 5 days ago

dced42f5

Reverts 7015be361f38e7149a3a4f97df81605d7e108501

gnecula committed 5 days ago

55f0a217

Split `fused_attention_stablehlo_test` into single-GPU and multi-GPU tests.

bchetioui committed 5 days ago

982ce8e0

Reverts 4e50bebcf74ca8049d27467ea05f4c5c14075afd

Google-ML-Automation committed 5 days ago

43892f52

Older