ROCm/jax

Pull Requests Commits

Ensure tests start with a clean MIOpen cache state

psanal35 committed 5 hours ago

ad4a5786

Resolve ROCm workflow conflict markers

psanal35 committed 8 hours ago

88e42dcc

[ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images

gulsumgudukbay committed 8 hours ago

ea0ac83f

[ROCm] Embed TheRock-aware wheel RUNPATH at link time

gulsumgudukbay committed 8 hours ago

15961355

[ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices

mminutoli committed 8 hours ago

90a9455f

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

bchetioui committed 9 hours ago

2fd83c70

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

bchetioui committed 11 hours ago

113cb4d2

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

gnecula committed 13 hours ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 13 hours ago

c24ca37d

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

hawkinsp committed 13 hours ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 14 hours ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 14 hours ago

59eea209

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

hawkinsp committed 14 hours ago

86a46a6e

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

PatriosTheGreat committed 15 hours ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 15 hours ago

2e4d7f65

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

hawkinsp committed 15 hours ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 15 hours ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 16 hours ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 16 hours ago

31c559c7

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

hawkinsp committed 16 hours ago

5ab345f1

Try to resolve test failure on ROCM.

akuegel committed 18 hours ago

305df9b6

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in Pallas

gnecula committed 18 hours ago

dc575083

Reverts dced42f5da959fad3268cbeedc024e0ffd8951a5

bchetioui committed 20 hours ago

452abf87

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

Google-ML-Automation committed 21 hours 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 1 day ago

05127a4c

Remove stale version guards after JAX v0.10.2 release.

hawkinsp committed 1 day ago

55d755f4

[JAX] Add `py_pjrt_host_callback` extension to jaxlib

Google-ML-Automation committed 1 day ago

cf6655fd

Merge pull request #38550 from hawkinsp:postrelease

Google-ML-Automation committed 1 day ago

6d4aff5f

Updates after JAX release 0.10.2

hawkinsp committed 1 day ago

8e784c54

Merge pull request #38538 from hawkinsp:continuoustest

Google-ML-Automation committed 1 day ago

c63bb6c1

Older