ROCm/jax

Pull Requests Commits

tsrw2048 committed 2 hours ago

a0bf6c11

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

mminutoli committed 7 days ago

c26569fb

[ROCm] Trim release validation matrix to TheRock only

magaonka-amd committed 8 days ago

73592795

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

magaonka-amd committed 8 days ago

4e057a47

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

magaonka-amd committed 9 days ago

fef5c3cd

Update ROCm runner labels to new labels

psanal35 committed 9 days ago

072e4511

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

magaonka-amd committed 9 days ago

c6572bd8

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 9 days ago

6bfbb491

Cherry pick env var leak fix from open upstream.

tsrw2048 committed 9 days ago

bd117038

Disable hipblaslt in TheRock to avoid current Tensile issues.

tsrw2048 committed 9 days ago

4df22d98

Add all specific target archs for ROCm

tsrw2048 committed 9 days ago

479be759

Trim down testing matrix

tsrw2048 committed 9 days ago

2ad6ebf6

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

tsrw2048 committed 9 days ago

c2dc41d1

Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 9 days ago

8501437c

Disable gpu_pallas_distributed_nvshmem_test_gpu which was failing at the

hawkinsp committed 15 days ago

990e6a0b

Synchronize PGLE profiler tests before stopping traces

zhenying-liu committed 15 days ago

4437b23d

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

hawkinsp committed 15 days ago

3815e4ea

Refactor Windows Docker GCE metadata access

nitins17 committed 15 days ago

272e3783

Prepare for JAX release 0.10.2

hawkinsp committed 16 days ago

310278eb

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

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

5597e903

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

WillFroom committed 18 days ago

728e4bfe

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

Google-ML-Automation committed 18 days ago

26113fa5

Pallas emit_pipeline raises ValueError on empty grid

vfdev-5 committed 18 days ago

c15bacc8

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

Google-ML-Automation committed 18 days ago

31a7aaba

Be more defensive when converting 3.13-ft to 3.13t in build scripts.

hawkinsp committed 19 days ago

8951c9be

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

Google-ML-Automation committed 19 days ago

5d4ec93e

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

Google-ML-Automation committed 20 days ago

242ad12b

Add mask annotation

Google-ML-Automation committed 20 days ago

31070f60

Merge pull request #38451 from jax-ml:effects-by-var

Google-ML-Automation committed 20 days ago

4dd5760e

Older