ROCm/jax

Pull Requests Commits

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

mminutoli committed 127 days ago

c2179c74

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

gulsumgudukbay committed 127 days ago

b0ab8909

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

gulsumgudukbay committed 127 days ago

d430e3bb

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

mminutoli committed 127 days ago

dfb1a195

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

yashk2810 committed 127 days ago

e441101c

[Mosaic GPU] Support swizzle in `is_valid_assignment`

golechwierowicz committed 127 days ago

d1d1cd25

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

Google-ML-Automation committed 128 days ago

786079b7

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

gnecula committed 128 days ago

b040143f

Make sure cosmetic changes to async collectives code

yashk2810 committed 128 days ago

fbbc7884

[Pallas/Mosaic GPU] Delete debug artefact.

bchetioui committed 128 days ago

0d88e2d3

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

bchetioui committed 128 days ago

2fd83c70

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

bchetioui committed 128 days ago

113cb4d2

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

gnecula committed 128 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 128 days ago

c24ca37d

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

hawkinsp committed 128 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 129 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 129 days ago

59eea209

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

hawkinsp committed 129 days ago

86a46a6e

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

PatriosTheGreat committed 129 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 129 days ago

2e4d7f65

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

hawkinsp committed 129 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 129 days ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 129 days ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 129 days ago

31c559c7

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

hawkinsp committed 129 days ago

5ab345f1

Try to resolve test failure on ROCM.

akuegel committed 129 days ago

305df9b6

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in Pallas

gnecula committed 129 days ago

dc575083

Reverts dced42f5da959fad3268cbeedc024e0ffd8951a5

bchetioui committed 129 days ago

452abf87

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

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

05127a4c

Older