ROCm/jax

Pull Requests Commits

Temporarily add unset JAX_PLATFORMS

tsrw2048 committed 52 days ago

66918cf7

Fix Rock path that was reverted.

tsrw2048 committed 54 days ago

7a106978

Add safeguards around testing script (timeouts and restarts).

tsrw2048 committed 54 days ago

98e4bc9c

Temporarily restrict testing to a single GPU

tsrw2048 committed 54 days ago

facef18a

Modify run_pytest_rocm.sh to just run single GPU tests sequentially.

tsrw2048 committed 54 days ago

ed8c391d

Add single and multi gpu versions of testing scripts.

tsrw2048 committed 55 days ago

a156a5cb

tsrw2048 committed 59 days ago

a0bf6c11

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

mminutoli committed 66 days ago

c26569fb

[ROCm] Trim release validation matrix to TheRock only

magaonka-amd committed 66 days ago

73592795

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

magaonka-amd committed 67 days ago

4e057a47

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

magaonka-amd committed 68 days ago

fef5c3cd

Update ROCm runner labels to new labels

psanal35 committed 68 days ago

072e4511

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

magaonka-amd committed 68 days ago

c6572bd8

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 68 days ago

6bfbb491

Cherry pick env var leak fix from open upstream.

tsrw2048 committed 68 days ago

bd117038

Disable hipblaslt in TheRock to avoid current Tensile issues.

tsrw2048 committed 68 days ago

4df22d98

Add all specific target archs for ROCm

tsrw2048 committed 68 days ago

479be759

Trim down testing matrix

tsrw2048 committed 68 days ago

2ad6ebf6

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

tsrw2048 committed 68 days ago

c2dc41d1

Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 68 days ago

8501437c

Disable gpu_pallas_distributed_nvshmem_test_gpu which was failing at the

hawkinsp committed 74 days ago

990e6a0b

Synchronize PGLE profiler tests before stopping traces

zhenying-liu committed 74 days ago

4437b23d

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

hawkinsp committed 74 days ago

3815e4ea

Refactor Windows Docker GCE metadata access

nitins17 committed 74 days ago

272e3783

Prepare for JAX release 0.10.2

hawkinsp committed 74 days ago

310278eb

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

allanrenucci committed 76 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 76 days ago

5597e903

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

WillFroom committed 76 days ago

728e4bfe

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

Google-ML-Automation committed 77 days ago

26113fa5

Pallas emit_pipeline raises ValueError on empty grid

vfdev-5 committed 77 days ago

c15bacc8

Older