ROCm/jax

Pull Requests Commits

Add AMD_SERIALIZE_KERNEL=3 for debug.

mminutoli committed 173 days ago

3253e03b

Cap the number of processes.

mminutoli committed 173 days ago

6d1a09e6

Remove 20 failures limit on pytest call.

mminutoli committed 173 days ago

d1a98da8

Remove timeout.

mminutoli committed 174 days ago

6acac54c

Strip down workflow to execute only rocm piece.

mminutoli committed 176 days ago

9518f1df

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 176 days ago

e8ab9de1

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 176 days ago

62d1b791

Add download-jax-rocm-wheels composite action

mminutoli committed 176 days ago

dafd2086

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 176 days ago

159db84c

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 176 days ago

47acadb8

adding run_pytest_rocm.sh

gulsumgudukbay committed 177 days ago

d1103ffc

Merge pull request #34063 from ezhulenev:nccl-version-flag

Google-ML-Automation committed 177 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 177 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

Google-ML-Automation committed 177 days ago

9bfb3764

[JAX] Enable `PjRtClient::CopyArraysForCrossHost` and `jax.device_put` to handle arrays whose shards require a mixture of host-local and cross-host transfers.

emilyfertig committed 177 days ago

1f30ced4

[pallas:sc] A few minor formatting tweaks in the lowering

superbobry committed 177 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 177 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 177 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 177 days ago

acc95f10

[pallas:mosaic] Nuked libtpu blocks older than 6 months

superbobry committed 177 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 177 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 177 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 177 days ago

e44e88d1

Raise an error if an unmapped argument to vmap is sharded on mesh axis that we are vmapping over

yashk2810 committed 177 days ago

3eecb6a7

Raise a better error message if an empty mesh is passed to shard_map

yashk2810 committed 177 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 177 days ago

25eebba9

Disable `tests/multiprocess:socket_transfer_test` since it's failing flakily in all backends.

emilyfertig committed 177 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 177 days ago

19833088

PR #34434: scaled_matmul_stablehlo_test: do not assert ordering within HLO

olupton committed 177 days ago

a5d92a06

Increase shard count for TPU and GPU tests back to 5 for api_test.py.

belitskiy committed 177 days ago

6041a2a9

Older