ROCm/jax

Pull Requests Commits

[ROCm] Added skip for "test_prim_tridiagonal_solve" tests on ROCm.

tsrw2048 committed 116 days ago

5eb5a215

Merge pull request #34643 from mattjj:nodceeffect-hasheq

Google-ML-Automation committed 117 days ago

0d66c531

add __hash__ / __eq__ to NoDCEEffect for serialization.py

mattjj committed 117 days ago

778ec599

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

Google-ML-Automation committed 117 days ago

d3259a54

_pathways.pyi: Allow creating an XLA CPU client.

Google-ML-Automation committed 117 days ago

18a983dc

[Mosaic GPU] Move kernel compilation out of the first execution.

allanrenucci committed 117 days ago

66e20704

[Pallas:SC] Fix tpu_sparsecore_pallas_test to work with TPUv7x

apaszke committed 117 days ago

4f32df24

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

danielsuo committed 117 days ago

0d34e25a

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

apaszke committed 117 days ago

2b518cc5

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

golechwierowicz committed 117 days ago

d4f0ec3f

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

golechwierowicz committed 117 days ago

a0aedf10

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

golechwierowicz committed 117 days ago

3b783f2d

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

bchetioui committed 117 days ago

882266e4

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

thcmbs committed 117 days ago

10993912

Reverts 0e56b42d9a163f67623e3bdc613d0230df0e900e

golechwierowicz committed 117 days ago

a21715c3

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

bchetioui committed 117 days ago

05201667

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

Google-ML-Automation committed 118 days ago

b19be1aa

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

danielsuo committed 118 days ago

6e671774

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

danielsuo committed 118 days ago

31c17496

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

golechwierowicz committed 118 days ago

0e56b42d

Automated Code Change

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

a1224782

Delete the //jax:experimental bazel target.

hawkinsp committed 118 days ago

125889e8

[Mosaic GPU] Use existing CUDA compute capability.

allanrenucci committed 119 days ago

58c0f572

Move array_garbage_collection_guard logic from PyArray_tp_clear to PyArray_tp_finalize.

ICGog committed 119 days ago

ad97537b

Update PGLE test to not expect spurious compilation cache misses.

Google-ML-Automation committed 120 days ago

f00e53fc

Disable `pytorch_interoperability_test` on GPU temporarily.

bchetioui committed 120 days ago

3648d95f

Enable pytorch_interoperability_test on GPU

ddelgadovargas-cyber committed 120 days ago

d8f3b88d

Add new Windows RBE configs

nitins17 committed 120 days ago

363082f9

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

Google-ML-Automation committed 120 days ago

7a5431c5

Older