ROCm/jax

Pull Requests Commits

Move the XLA pin to the head of rocm-jaxlib-v0.10.2

gulsumgudukbay committed 76 days ago

6f7ac687

[ROCm] Fix ROCm version in plugin wheel metadata (#844)

gulsumgudukbay committed 79 days ago

Verified 20ca79de

[ROCm] Simplify pytest parallelism calculation for ROCm CI

magaonka-amd committed 82 days ago

7f98f407

[ROCm] Split the multi-accelerator tests into their own script (#833)

gulsumgudukbay committed 82 days ago

Verified cae8aad3

[ROCm] Drop the rocm-sdk PATH block from run_pytest_rocm.sh

gulsumgudukbay committed 88 days ago

e5c733d3

Derive the ROCm plugin package name from the ROCm major version

gulsumgudukbay committed 94 days ago

4633d1ef

[ROCm][Pallas] Skip fused attention fwd tests that exceed device resources - backport v0.10.2 (#819)

tsrw2048 committed 103 days ago

Verified 0dad77d6

[ROCm] Add release pipeline CI workflows (#808)

mminutoli committed 129 days ago

c26569fb

[ROCm] Trim release validation matrix to TheRock only

magaonka-amd committed 129 days ago

73592795

[ROCm] Bump pinned XLA to rocm-jaxlib-v0.10.2 HEAD (ROCm/xla#993)

magaonka-amd committed 130 days ago

4e057a47

[ROCm] Use new gfx950 runner labels in release validation workflow

magaonka-amd committed 131 days ago

fef5c3cd

Update ROCm runner labels to new labels

psanal35 committed 131 days ago

072e4511

[ROCm] Point pinned XLA at ROCm/xla fork for the 0.10.2 release

magaonka-amd committed 131 days ago

c6572bd8

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 131 days ago

6bfbb491

Cherry pick env var leak fix from open upstream.

tsrw2048 committed 131 days ago

bd117038

Disable hipblaslt in TheRock to avoid current Tensile issues.

tsrw2048 committed 131 days ago

4df22d98

Add all specific target archs for ROCm

tsrw2048 committed 131 days ago

479be759

Trim down testing matrix

tsrw2048 committed 131 days ago

2ad6ebf6

Add HSA_NO_SCRATCH_RECLAIM env var and update Rock images to 7.14 in release workflow

tsrw2048 committed 131 days ago

c2dc41d1

Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 131 days ago

8501437c

Disable gpu_pallas_distributed_nvshmem_test_gpu which was failing at the

hawkinsp committed 137 days ago

990e6a0b

Synchronize PGLE profiler tests before stopping traces

zhenying-liu committed 137 days ago

4437b23d

Add a workaround for linux aarch64 python 3.14t pytest builds.

hawkinsp committed 137 days ago

3815e4ea

Refactor Windows Docker GCE metadata access

nitins17 committed 137 days ago

272e3783

Prepare for JAX release 0.10.2

hawkinsp committed 137 days ago

310278eb

[Pallas:MGPU] Add support for batch dimensions to `async_store_tmem`.

allanrenucci committed 139 days ago

7db3a70e

[Mosaic:GPU] Disable set_xla_gpu_enable_allocator_spatial_partitioning for nvshmem test case, since nvshmem can't be used within the same allocation.

PatriosTheGreat committed 139 days ago

5597e903

Add functionality to blanket disable the custom call safety checks in export.

WillFroom committed 139 days ago

728e4bfe

Merge pull request #38432 from vfdev-5:pallas-emit-pipeline-raises-error-empty-grid

Google-ML-Automation committed 139 days ago

26113fa5

Pallas emit_pipeline raises ValueError on empty grid

vfdev-5 committed 140 days ago

c15bacc8

Older