ROCm/jax

Pull Requests Commits

Skip test_batch_axis_sharding_jvp13 on ROCm due to rocSolver issue

AratiGanesh committed 22 days ago

10093c64

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

Google-ML-Automation committed 48 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 48 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 48 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 48 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 48 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 48 days ago

acc95f10

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

superbobry committed 48 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 48 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 48 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 48 days ago

e44e88d1

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

yashk2810 committed 48 days ago

3eecb6a7

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

yashk2810 committed 49 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 49 days ago

25eebba9

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

emilyfertig committed 49 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 49 days ago

19833088

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

olupton committed 49 days ago

a5d92a06

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

belitskiy committed 49 days ago

6041a2a9

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

superbobry committed 49 days ago

9c25bbc2

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

apaszke committed 49 days ago

0cd5ca77

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

apaszke committed 49 days ago

51d35912

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

Google-ML-Automation committed 49 days ago

7a5e13d7

Add Fuser support for jnp.tile input fusions

Google-ML-Automation committed 49 days ago

156f7d2e

Bump XLA version.

mwhittaker committed 49 days ago

cdfe904a

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

ybaturina committed 49 days ago

a84ca0c5

[Pallas] Create trace_value primitive for dynamic value logging

Google-ML-Automation committed 49 days ago

02939fef

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

Google-ML-Automation committed 49 days ago

6a9f3d98

Follow up pr after triton integration cl and tokamax pr.

loislo committed 49 days ago

91223b93

Older