ROCm/jax

Pull Requests Commits

gulsumgudukbay committed 31 days ago

70e9b791

mark mosaic tests as xfail due to mosaic not being supported on ROCm

gulsumgudukbay committed 31 days ago

91244f71

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 31 days ago

15199269

[jax] bound traceback of pallas_call to within the kernel.

cezheng committed 31 days ago

fe359a15

Clean up a few TODOs in pjit after FlatTree change

yashk2810 committed 31 days ago

aeccdff4

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

Google-ML-Automation committed 31 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 31 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 31 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 31 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 31 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 31 days ago

acc95f10

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

superbobry committed 31 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 31 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 31 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 31 days ago

e44e88d1

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

yashk2810 committed 31 days ago

3eecb6a7

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

yashk2810 committed 31 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 31 days ago

25eebba9

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

emilyfertig committed 31 days ago

5b2e5673

Fix wheel sources tests for Windows platform.

ybaturina committed 31 days ago

19833088

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

olupton committed 31 days ago

a5d92a06

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

belitskiy committed 31 days ago

6041a2a9

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

superbobry committed 32 days ago

9c25bbc2

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

apaszke committed 32 days ago

0cd5ca77

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

apaszke committed 32 days ago

51d35912

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

Google-ML-Automation committed 32 days ago

7a5e13d7

Add Fuser support for jnp.tile input fusions

Google-ML-Automation committed 32 days ago

156f7d2e

Bump XLA version.

mwhittaker committed 32 days ago

cdfe904a

Older