ROCm/jax

Pull Requests Commits

Enabled testTridiagonal to run on ROCm devices as relevant tests are all passing.

tsrw2048 committed 45 days ago

5aa84b49

[Mosaic GPU] Allow load/store_tiled with refs where the tiling rank is smaller than the logical rank

apaszke committed 45 days ago

cb24d7fc

[XLA:MGPU] Port warp_indices and lane_indices to C++.

golechwierowicz committed 45 days ago

0cf96f75

Merge pull request #34402 from gnecula:clean_dynamic_shape_replacement

Google-ML-Automation committed 45 days ago

64301806

Merge pull request #34425 from andportnoy:aportnoy/mosaic-gpu-fence-proxy-async

Google-ML-Automation committed 45 days ago

c32e170b

[pallas] Minor cleanup of typing annotations for dynamic_shape_replacement_fn

gnecula committed 45 days ago

a51d661b

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

Google-ML-Automation committed 45 days ago

7c5c109d

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

Google-ML-Automation committed 46 days ago

ff99faf2

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

Google-ML-Automation committed 47 days ago

c6568036

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 47 days ago

15199269

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

cezheng committed 47 days ago

fe359a15

Clean up a few TODOs in pjit after FlatTree change

yashk2810 committed 47 days ago

aeccdff4

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

Google-ML-Automation committed 47 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 47 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 47 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 47 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 47 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 47 days ago

acc95f10

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

superbobry committed 47 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 47 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 47 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 47 days ago

e44e88d1

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

yashk2810 committed 47 days ago

3eecb6a7

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

yashk2810 committed 47 days ago

d927d96b

Re-enable `socket_transfer_test` internally.

emilyfertig committed 47 days ago

25eebba9

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

emilyfertig committed 47 days ago

5b2e5673

[Mosaic GPU] Fix MosaicGpuDialectTest::test_wgmma_row_col_store

andportnoy committed 47 days ago

f983f4bf

Fix wheel sources tests for Windows platform.

ybaturina committed 47 days ago

19833088

Older