ROCm/jax

Pull Requests Commits

[Pallas:SC] Fix tpu_sparsecore_pallas_test to work with TPUv7x

apaszke committed 125 days ago

4f32df24

[pmap] Expose a no data move/copy sharding metadata rewrap.

danielsuo committed 125 days ago

0d34e25a

[Pallas:SC] Disable a failing debug check test on TPUv7x

apaszke committed 126 days ago

2b518cc5

[XLA:MGPU] Port utils.memref_transpose to C++

golechwierowicz committed 126 days ago

d4f0ec3f

[XLA:MGPU] Port utils.memref_slice to C++.

golechwierowicz committed 126 days ago

a0aedf10

[XLA:MGPU] Port utils.memref_unfold to C++

golechwierowicz committed 126 days ago

3b783f2d

[Pallas/Mosaic GPU] Disable `test_print_wgmma_tiled_layout` on B200.

bchetioui committed 126 days ago

882266e4

[Mosaic GPU] Fix use-after-free in kernel cache by ensuring pointer stability.

thcmbs committed 126 days ago

10993912

Reverts 0e56b42d9a163f67623e3bdc613d0230df0e900e

golechwierowicz committed 126 days ago

a21715c3

[Pallas/Mosaic GPU] Add `nvdisasm` dependency to `mosaic_gpu_test`.

bchetioui committed 126 days ago

05201667

Fixed few typos in the training cookbook

vfdev-5 committed 126 days ago

a8470ce4

Move collective metadata initialization to the Initialize stage of the Mosaic GPU custom call.

Google-ML-Automation committed 126 days ago

b19be1aa

[pmap] Don't check devices on every invocation of a pmapped function.

danielsuo committed 126 days ago

6e671774

[pmap] Don't recreate the wrapped function on each invocation of pmap.

danielsuo committed 126 days ago

31c17496

[XLA:MGPU] Port the entire TiledLayout class to C++.

golechwierowicz committed 126 days ago

0e56b42d

Automated Code Change

Google-ML-Automation committed 126 days ago

0a46c91f

Add a test to make sure axis_index behaves as expected under vmap when axis_data.name is not in axes.

yashk2810 committed 126 days ago

a1224782

Delete the //jax:experimental bazel target.

hawkinsp committed 126 days ago

125889e8

[Mosaic GPU] Use existing CUDA compute capability.

allanrenucci committed 127 days ago

58c0f572

Move array_garbage_collection_guard logic from PyArray_tp_clear to PyArray_tp_finalize.

ICGog committed 127 days ago

ad97537b

Update PGLE test to not expect spurious compilation cache misses.

Google-ML-Automation committed 128 days ago

f00e53fc

Disable `pytorch_interoperability_test` on GPU temporarily.

bchetioui committed 128 days ago

3648d95f

Enable pytorch_interoperability_test on GPU

ddelgadovargas-cyber committed 128 days ago

d8f3b88d

Add new Windows RBE configs

nitins17 committed 128 days ago

363082f9

Do not cache executables when MeshComputations.compile is called with a device list.

Google-ML-Automation committed 128 days ago

7a5431c5

Make compile methods in `xla::ifrt::Compiler` return a future of an executable

junwhanahn committed 128 days ago

3d1ea0b2

Introduce `jax.extend.clear_in_memory_compilation_cache` to allow users to clear the LRU cache for mesh executables during compilation

krishnaharidasan committed 128 days ago

7fedc0e5

Support cross-host `device_put` for a few new cases.

emilyfertig committed 128 days ago

bdcc93f7

Disable pytorch_interoperability_test on GPU since it's broken

yashk2810 committed 128 days ago

3cf7ece1

Fix error checking logic in check_unreduced_args to only error if axes of a primitive and aval's unreduced intersect.

yashk2810 committed 128 days ago

0b986900

Newer Older