ROCm/jax

Pull Requests Commits

Strip down workflow to execute only rocm piece.

mminutoli committed 149 days ago

10033636

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 149 days ago

76386e88

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 149 days ago

e17c856e

Add download-jax-rocm-wheels composite action

mminutoli committed 149 days ago

85ffa3ca

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 149 days ago

5f13602e

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 149 days ago

b116d57b

adding run_pytest_rocm.sh

gulsumgudukbay committed 154 days ago

d1103ffc

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

Google-ML-Automation committed 154 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 154 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 154 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 155 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 155 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 155 days ago

acc95f10

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

superbobry committed 155 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 155 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 155 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 155 days ago

e44e88d1

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

yashk2810 committed 155 days ago

3eecb6a7

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

yashk2810 committed 155 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 155 days ago

25eebba9

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

emilyfertig committed 155 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 155 days ago

19833088

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

olupton committed 155 days ago

a5d92a06

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

belitskiy committed 155 days ago

6041a2a9

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

superbobry committed 155 days ago

9c25bbc2

[Mosaic GPU] Lock down 8-bit pointwise ops (except bitwise and cmp) + add lshift/rshift

apaszke committed 155 days ago

0cd5ca77

[Mosaic GPU] Reenable the use of redux.sync, but only when all lanes participate

apaszke committed 155 days ago

51d35912

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

Google-ML-Automation committed 155 days ago

7a5e13d7

Older