ROCm/jax

Pull Requests Commits

Temp: Sweep per-GPU worker count to find the best value

psanal35 committed 114 days ago

63140c7c

Cap xdist workers per GPU to avoid oversubscription

psanal35 committed 114 days ago

2327338d

Isolate MIOpen cache per xdist worker to avoid lock contention

psanal35 committed 114 days ago

b7935f2f

Resolve ROCm workflow conflict markers

psanal35 committed 114 days ago

f4b6c808

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

gulsumgudukbay committed 114 days ago

4da85913

[ROCm] Update CI runner labels to amd-do-linux.jax.gpu.gfx950.{1,4,8} (#807)

mminutoli committed 114 days ago

Verified cb097f7a

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

gulsumgudukbay committed 114 days ago

b0ab8909

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

gulsumgudukbay committed 114 days ago

d430e3bb

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

mminutoli committed 114 days ago

dfb1a195

Handle singleton dimensions in reshape sharding rule in a more robust way.

yashk2810 committed 114 days ago

e441101c

[Mosaic GPU] Support swizzle in `is_valid_assignment`

golechwierowicz committed 114 days ago

d1d1cd25

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

Google-ML-Automation committed 115 days ago

786079b7

[trace_to_jaxpr] Port ops and TPU Pallas tests to use trace_to_jaxpr

gnecula committed 115 days ago

b040143f

Make sure cosmetic changes to async collectives code

yashk2810 committed 115 days ago

fbbc7884

[Pallas/Mosaic GPU] Delete debug artefact.

bchetioui committed 115 days ago

0d88e2d3

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

bchetioui committed 115 days ago

2fd83c70

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

bchetioui committed 115 days ago

113cb4d2

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

gnecula committed 115 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 115 days ago

c24ca37d

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

hawkinsp committed 115 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 115 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 115 days ago

59eea209

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

hawkinsp committed 115 days ago

86a46a6e

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

PatriosTheGreat committed 115 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 115 days ago

2e4d7f65

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

hawkinsp committed 115 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 115 days ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 115 days ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 115 days ago

31c559c7

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

hawkinsp committed 115 days ago

5ab345f1

Older