ROCm/jax

Pull Requests Commits

[Mosaic GPU][NFC] Pass `CompiledKernel` by pointer to `CachedInit`.

allanrenucci committed 26 days ago

745cbe2f

[Pallas/Mosaic GPU] Use `self.capture_stdout` instead of `jtu.capture_stdout` in `mosaic_gpu_test`.

bchetioui committed 26 days ago

7241ee09

Fix XLA integrity in MODULE.bazel

meteorcloudy committed 26 days ago

9663c503

[Pallas] Disable a GPU interpreter test that deadlocks sometimes

apaszke committed 26 days ago

5cab99f9

[Mosaic GPU] Slightly refactor the argument idx search routine

apaszke committed 27 days ago

f0693d90

[Mosaic GPU] Generalize the fast i8 -> bf16 conversion to allow more vector lengths

apaszke committed 27 days ago

3b5d7209

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

Google-ML-Automation committed 27 days ago

de7e49e1

Merge pull request #34610 from magaonka-amd:skip-sparse-hipsparse-tests-upstream

Google-ML-Automation committed 27 days ago

cf73ceae

Correctly check dtype equality in `is_matmul_supported`.

Google-ML-Automation committed 27 days ago

979b537a

Support `device_put` of an uncommitted array to a single global device in McJAX.

emilyfertig committed 27 days ago

40b9554d

Reverts d3259a54160ee3a3724a17dce143d3212e5454b9

nitins17 committed 27 days ago

e0b62e82

Merge pull request #34632 from vfdev-5:fixd-typos-in-training-cookbook-

Google-ML-Automation committed 27 days ago

9f243da9

Merge pull request #34643 from mattjj:nodceeffect-hasheq

Google-ML-Automation committed 27 days ago

0d66c531

add __hash__ / __eq__ to NoDCEEffect for serialization.py

mattjj committed 27 days ago

778ec599

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

Google-ML-Automation committed 27 days ago

d3259a54

_pathways.pyi: Allow creating an XLA CPU client.

Google-ML-Automation committed 27 days ago

18a983dc

Skip sparse tests on ROCm due to hipSPARSE issue

magaonka-amd committed 27 days ago

3acc4108

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

allanrenucci committed 27 days ago

66e20704

[Pallas:SC] Fix tpu_sparsecore_pallas_test to work with TPUv7x

apaszke committed 27 days ago

4f32df24

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

danielsuo committed 27 days ago

0d34e25a

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

apaszke committed 27 days ago

2b518cc5

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

golechwierowicz committed 27 days ago

d4f0ec3f

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

golechwierowicz committed 27 days ago

a0aedf10

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

golechwierowicz committed 28 days ago

3b783f2d

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

bchetioui committed 28 days ago

882266e4

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

thcmbs committed 28 days ago

10993912

Reverts 0e56b42d9a163f67623e3bdc613d0230df0e900e

golechwierowicz committed 28 days ago

a21715c3

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

bchetioui committed 28 days ago

05201667

Fixed few typos in the training cookbook

vfdev-5 committed 28 days ago

a8470ce4

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

Google-ML-Automation committed 28 days ago

b19be1aa

Newer Older