ROCm/jax

Pull Requests Commits

Remove PYTHONSAFEPATH

mminutoli committed 179 days ago

42e7523c

fixup! Add download-jax-rocm-wheels composite action

mminutoli committed 179 days ago

290d71ea

Strip down workflow to execute only rocm piece.

mminutoli committed 179 days ago

2523188f

fixup! Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 179 days ago

2106ff8c

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 180 days ago

a99605b9

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 180 days ago

b3ebfd48

Add download-jax-rocm-wheels composite action

mminutoli committed 180 days ago

d23367a9

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 180 days ago

b2340e2e

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 180 days ago

0ad3b479

adding run_pytest_rocm.sh

gulsumgudukbay committed 180 days ago

d1103ffc

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

Google-ML-Automation committed 180 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 180 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 180 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 180 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 180 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 180 days ago

acc95f10

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

superbobry committed 180 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 180 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 180 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 180 days ago

e44e88d1

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

yashk2810 committed 180 days ago

3eecb6a7

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

yashk2810 committed 180 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 180 days ago

25eebba9

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

emilyfertig committed 180 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 180 days ago

19833088

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

olupton committed 180 days ago

a5d92a06

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

belitskiy committed 180 days ago

6041a2a9

[mosaic] Propagate the layout from tpu.erase_memref_layout operand to tpu.log_buffer

superbobry committed 180 days ago

9c25bbc2

Older