ROCm/jax

Pull Requests Commits

Took suggestion from PR to merge CUDA guards for readability.

tsrw2048 committed 297 days ago

1aa82e36

Added support for GESVDJ on ROCm GPUs.

tsrw2048 committed 299 days ago

5957253b

Merge pull request #34288 from olupton:smem-limit

Google-ML-Automation committed 299 days ago

1033308a

Artificially limit SMEM capacity for tests

olupton committed 299 days ago

490b6472

[XLA:MGPU] Port registers_shape, shape_from_registers_shape, registers_element_type methods.

golechwierowicz committed 299 days ago

ee387893

[Pallas/Mosaic GPU] Expose `plgpu.wgmma_accumulator_load`.

bchetioui committed 299 days ago

08a8bf71

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

apaszke committed 300 days ago

cb24d7fc

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

golechwierowicz committed 300 days ago

0cf96f75

Merge pull request #34402 from gnecula:clean_dynamic_shape_replacement

Google-ML-Automation committed 300 days ago

64301806

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

Google-ML-Automation committed 300 days ago

c32e170b

[pallas] Minor cleanup of typing annotations for dynamic_shape_replacement_fn

gnecula committed 300 days ago

a51d661b

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

Google-ML-Automation committed 300 days ago

7c5c109d

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

Google-ML-Automation committed 301 days ago

ff99faf2

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

Google-ML-Automation committed 302 days ago

c6568036

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 302 days ago

15199269

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

cezheng committed 302 days ago

fe359a15

Clean up a few TODOs in pjit after FlatTree change

yashk2810 committed 302 days ago

aeccdff4

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

Google-ML-Automation committed 302 days ago

cb2d81de

Add support for specifying NCCL version with nccl_version flag

ezhulenev committed 302 days ago

189845a8

Merge pull request #33186 from ROCm:phambinh/ci_fix_build_nvidia

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

1f30ced4

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

superbobry committed 302 days ago

959ce9f5

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 302 days ago

58c64b98

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 302 days ago

b685442a

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 302 days ago

acc95f10

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

superbobry committed 302 days ago

0353959c

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 302 days ago

d0e8455a

Use arena allocation for protobufs when serializing HloModules.

hawkinsp committed 302 days ago

307a2553

Add libtpu date guard for failing test.

mwhittaker committed 302 days ago

e44e88d1

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

yashk2810 committed 302 days ago

3eecb6a7

Older