ROCm/jax

Pull Requests Commits

ScXfjiang committed 13 days ago

c5e3b96b

register rocm to block_scaled_dot lowering path

ScXfjiang committed 15 days ago

283e4c77

[Mosaic GPU] Add support for arbitrary reshapes of contiguous refs

apaszke committed 15 days ago

dd12a48e

[Mosaic GPU] Optimize the computation of tcgen05.mma matrix descriptors

apaszke committed 15 days ago

92621c77

[Mosaic GPU][NFC] Refactor: Use a macro for CUDA error checking.

allanrenucci committed 15 days ago

a83b4d68

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

Google-ML-Automation committed 15 days ago

167dc0df

Remove `testonly` attribute from `pallas_test_util` to include the target in the additional wheel dependencies.

ybaturina committed 15 days ago

fa27af7b

[PjRt-IFRT] Tighten `PjRtExecutable::Create()` to take an MLIR module

hyeontaek committed 15 days ago

6168f5cd

Update references deprecated JAX build targets

hawkinsp committed 15 days ago

71d9ec14

Reverts 2f7c9073f4978373e04289c91b7fdccf0f8c06c7

mwhittaker committed 15 days ago

7af0e643

Rename if_cloud_tpu_at_least to is_cloud_tpu_at_least

sharadmv committed 15 days ago

36cd74a1

Add `pallas_test_util.py` to the test runfiles as a follow-up after https://github.com/jax-ml/jax/pull/33704.

ybaturina committed 15 days ago

bf766101

AOT: remove explicit lojax stage (`Fallen`) and extend `Traced` instead

froystig committed 15 days ago

2e548e5b

Change PjRt to use new copy of coordination service.

mwhittaker committed 15 days ago

2f7c9073

Add test case to debug_print scalar in Pallas.

IvyZX committed 15 days ago

334ef7ef

Merge pull request #33048 from IvyZX:core_map

Google-ML-Automation committed 15 days ago

2e006a9b

Merge pull request #33735 from jakevdp:fix-np-24-signature

Google-ML-Automation committed 15 days ago

47a5c4b5

Add `example_libraries/__init__.py` to jax wheel sources as a follow-up after the cleanup in https://github.com/jax-ml/jax/pull/33723.

ybaturina committed 15 days ago

f9e9e1da

[test] make signature error test robust to NumPy version

jakevdp committed 15 days ago

766f1899

Merge pull request #33705 from jakevdp:del-jax-metal

Google-ML-Automation committed 15 days ago

c100eb66

Merge pull request #33537 from vfdev-5:fix-typo-docstring-make_array_from_process_local_data

Google-ML-Automation committed 15 days ago

9a4cf171

Merge pull request #33737 from mattjj:issue33689-2

Google-ML-Automation committed 15 days ago

081c1532

[Pallas] Skip interpret tests on accelerators

justinjfu committed 15 days ago

3ba9f282

add test for #33714

mattjj committed 15 days ago

42eafcb7

[Pallas] Hardcode einsum flops estimate.

justinjfu committed 15 days ago

a8bd2163

PR #33714: always apply stop_gradient_p, even to exact dtypes

mattjj committed 15 days ago

7f31fa31

Merge pull request #33711 from mattjj:issue30787

Google-ML-Automation committed 15 days ago

0b3986b6

Merge pull request #33698 from jakevdp:check-static-indices

Google-ML-Automation committed 15 days ago

ce9bcf80

Make in_avals and out_avals match as well as we can with in_shardings and out_shardings respectively during lowering.

yashk2810 committed 15 days ago

ca4d7884

Merge pull request #33720 from olupton:thor

Google-ML-Automation committed 15 days ago

0625fdc7

Older