ROCm/jax

Pull Requests Commits

Isolate MIOpen cache per xdist worker to avoid lock contention

psanal35 committed 119 days ago

b96c0a91

Resolve ROCm workflow conflict markers

psanal35 committed 119 days ago

88e42dcc

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

gulsumgudukbay committed 119 days ago

ea0ac83f

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

gulsumgudukbay committed 119 days ago

15961355

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

mminutoli committed 119 days ago

90a9455f

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

bchetioui committed 119 days ago

2fd83c70

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

bchetioui committed 119 days ago

113cb4d2

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

gnecula committed 119 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 119 days ago

c24ca37d

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

hawkinsp committed 119 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 119 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 119 days ago

59eea209

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

hawkinsp committed 119 days ago

86a46a6e

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

PatriosTheGreat committed 120 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 120 days ago

2e4d7f65

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

hawkinsp committed 120 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 120 days ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 120 days ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 120 days ago

31c559c7

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

hawkinsp committed 120 days ago

5ab345f1

Try to resolve test failure on ROCM.

akuegel committed 120 days ago

305df9b6

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in Pallas

gnecula committed 120 days ago

dc575083

Reverts dced42f5da959fad3268cbeedc024e0ffd8951a5

bchetioui committed 120 days ago

452abf87

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

Google-ML-Automation committed 120 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 120 days ago

05127a4c

Remove stale version guards after JAX v0.10.2 release.

hawkinsp committed 120 days ago

55d755f4

[JAX] Add `py_pjrt_host_callback` extension to jaxlib

Google-ML-Automation committed 120 days ago

cf6655fd

Merge pull request #38550 from hawkinsp:postrelease

Google-ML-Automation committed 120 days ago

6d4aff5f

Updates after JAX release 0.10.2

hawkinsp committed 120 days ago

8e784c54

Merge pull request #38538 from hawkinsp:continuoustest

Google-ML-Automation committed 120 days ago

c63bb6c1

Older