ROCm/jax

Pull Requests Commits

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

shurale-nkn committed 12 hours ago

0c672603

[Mosaic] Prove divisibility of loop induction vars

naummo committed 15 hours ago

3ad2cd96

Merge pull request #35965 from gnecula:export_doc

Google-ML-Automation committed 15 hours ago

bc9bbe70

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

gnecula committed 15 hours ago

d48994d7

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

Google-ML-Automation committed 17 hours ago

7c5f74c8

Merge pull request #35959 from olupton:cta_group_2

Google-ML-Automation committed 17 hours ago

1dc08cf5

Disable asan for lax_numpy_hijax_test on TPU

Jake VanderPlas committed 17 hours ago

a21359d9

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

tlongeri committed 17 hours ago

d75c352b

Bump the maximum CUDA PJRT plugin wheel size.

akuegel committed 18 hours ago

18f1ed08

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

olupton committed 20 hours ago

cf3313c7

Merge pull request #35218 from zhenying-liu:input_output_alias

Google-ML-Automation committed 23 hours ago

f647bf2d

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

alekstheod committed 23 hours ago

ce621cf5

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

Google-ML-Automation committed 23 hours ago

724bc41b

Merge pull request #33383 from sergachev:fix_ag_test

Google-ML-Automation committed 23 hours ago

ecb523d2

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

yashk2810 committed 1 day ago

53a4d874

Rename `_identity_fn` to `_device_put_reshard`.

tomhennigan committed 1 day ago

4eb00c9b

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

Google-ML-Automation committed 1 day ago

e2022582

Fix convert_element_type unreduced propagation

yashk2810 committed 1 day ago

0422b1a7

tree_util.register_dataclass: fix drop_fields bug

jakevdp committed 1 day ago

f892f6bf

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

tomhennigan committed 1 day ago

20984312

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

Kevin Wang committed 1 day ago

ce5c3824

Merge pull request #35939 from jakevdp:del-arraylike

Google-ML-Automation committed 1 day ago

348bff56

[cleanup] remove unnecessary call to lax.asarray

jakevdp committed 1 day 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 1 day ago

02007235

[Mosaic GPU] Add async_copy_sparse_metadata_smem_to_tmem dialect lowering.

golechwierowicz committed 1 day ago

12a24495

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

golechwierowicz committed 1 day ago

555baca7

[Mosaic GPU][NFC] Cleanup error reporting.

allanrenucci committed 1 day ago

3af20779

[Mosaic GPU] Remove a_sparse_metadata availability check.

golechwierowicz committed 1 day ago

196f85b9

Fix a flake in LoggingTest.test_debug_logging.

hawkinsp committed 1 day ago

3c0df6e3

Disable hypothesis too_slow checks under ASAN/TSAN/MSAN.

hawkinsp committed 1 day ago

054eaacf

Older