ROCm/jax

Pull Requests Commits

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

Google-ML-Automation committed 21 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 21 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 21 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 21 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 21 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 21 days ago

acc95f10

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

superbobry committed 21 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 21 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 21 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 21 days ago

e44e88d1

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

yashk2810 committed 21 days ago

3eecb6a7

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

yashk2810 committed 21 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 21 days ago

25eebba9

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

emilyfertig committed 21 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 21 days ago

19833088

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

olupton committed 22 days ago

a5d92a06

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

belitskiy committed 22 days ago

6041a2a9

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

superbobry committed 22 days ago

9c25bbc2

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

apaszke committed 22 days ago

0cd5ca77

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

apaszke committed 22 days ago

51d35912

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

Google-ML-Automation committed 22 days ago

7a5e13d7

Add Fuser support for jnp.tile input fusions

Google-ML-Automation committed 22 days ago

156f7d2e

Bump XLA version.

mwhittaker committed 22 days ago

cdfe904a

Add `testonly=True` to py_import targets that depend on testonly `wheel_additives` target.

ybaturina committed 22 days ago

a84ca0c5

[Pallas] Create trace_value primitive for dynamic value logging

Google-ML-Automation committed 22 days ago

02939fef

Merge pull request #34418 from jakevdp:hijax-shard-map

Google-ML-Automation committed 22 days ago

6a9f3d98

Follow up pr after triton integration cl and tokamax pr.

loislo committed 22 days ago

91223b93

Migrate TSAN workflow to use RBE.

belitskiy committed 22 days ago

acdd012a

Older