ROCm/jax

Pull Requests Commits

Add ROCm support for scaled matrix multiplication and scaled dot lowering.

shurale-nkn committed 22 days ago

246853b1

[Mosaic] Prove divisibility of loop induction vars

naummo committed 28 days ago

3ad2cd96

Merge pull request #35965 from gnecula:export_doc

Google-ML-Automation committed 28 days ago

bc9bbe70

[export] Reflow the export.md file to keep lines under 80 char long

gnecula committed 28 days ago

d48994d7

Merge pull request #35932 from ROCm:rocm/s3-wheel-downloads

Google-ML-Automation committed 28 days ago

7c5f74c8

Merge pull request #35959 from olupton:cta_group_2

Google-ML-Automation committed 28 days ago

1dc08cf5

Disable asan for lax_numpy_hijax_test on TPU

Jake VanderPlas committed 28 days ago

a21359d9

[Mosaic] Enable internal fptoi/itofp ops for scalar types

tlongeri committed 28 days ago

d75c352b

Bump the maximum CUDA PJRT plugin wheel size.

akuegel committed 28 days ago

18f1ed08

[Mosaic GPU] Skip .cta_group::2 tests on non-tcgen05 hardware

olupton committed 28 days ago

cf3313c7

Merge pull request #35218 from zhenying-liu:input_output_alias

Google-ML-Automation committed 28 days ago

f647bf2d

PR #35516: [ROCm] Include prebuilt rocm plugin pkgs as deps to tests

alekstheod committed 28 days ago

ce621cf5

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

Google-ML-Automation committed 28 days ago

724bc41b

Merge pull request #33383 from sergachev:fix_ag_test

Google-ML-Automation committed 28 days ago

ecb523d2

Add a test to make sure manual -> unreduced in shardy works. Fixes https://github.com/jax-ml/jax/issues/35934

yashk2810 committed 28 days ago

53a4d874

Rename `_identity_fn` to `_device_put_reshard`.

tomhennigan committed 29 days ago

4eb00c9b

Merge pull request #35938 from jakevdp:fix-drop-fields

Google-ML-Automation committed 29 days ago

e2022582

Fix convert_element_type unreduced propagation

yashk2810 committed 29 days ago

0422b1a7

tree_util.register_dataclass: fix drop_fields bug

jakevdp committed 29 days ago

f892f6bf

Speed up computation of logical_device_ids in reshard (`O(N^2)` -> `O(N)`).

tomhennigan committed 29 days ago

20984312

Increase the number of args passed to the VJP too many args error message by 1.

Kevin Wang committed 29 days ago

ce5c3824

Merge pull request #35939 from jakevdp:del-arraylike

Google-ML-Automation committed 29 days ago

348bff56

[cleanup] remove unnecessary call to lax.asarray

jakevdp committed 29 days ago

53ee1d0e

Make sure to use `to_ct_spec()` when unsharding Zeros in shard_map_transpose. Also use `to_ct_aval` in _flatten_bwd in custom_vjp

yashk2810 committed 29 days ago

02007235

[Mosaic GPU] Add async_copy_sparse_metadata_smem_to_tmem dialect lowering.

golechwierowicz committed 29 days ago

12a24495

[Mosaic GPU] Test tcgen05_mma_sparse_metadata layout inference when LHS is in TMEM.

golechwierowicz committed 29 days ago

555baca7

[Mosaic GPU][NFC] Cleanup error reporting.

allanrenucci committed 29 days ago

3af20779

[Mosaic GPU] Remove a_sparse_metadata availability check.

golechwierowicz committed 29 days ago

196f85b9

Fix a flake in LoggingTest.test_debug_logging.

hawkinsp committed 29 days ago

3c0df6e3

Disable hypothesis too_slow checks under ASAN/TSAN/MSAN.

hawkinsp committed 29 days ago

054eaacf

Older