ROCm/jax

Pull Requests Commits

[Mosaic GPU] Improve tcgen05 preconditions error messaging.

golechwierowicz committed 28 days ago

93772b38

[NFC] Streamline `copy_gmem_to_smem_p` lowering registration.

allanrenucci committed 29 days ago

f6fd0352

[NFC] Streamline lowering registration for warp primitive semantics.

allanrenucci committed 29 days ago

96531f06

[Mosaic GPU] Add Pallas WG lowering for async_copy_sparse_metadata_to_tmem.

golechwierowicz committed 29 days ago

e4d133a8

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

Google-ML-Automation committed 29 days ago

464b047e

Replace `HloSharding::IsTileMaximal` with `HloSharding::IsReplicatedOrSingleDevice` since maximal is vague.

ZixuanJiang committed 29 days ago

3558a4af

Put `...` of overloads to newlines for some APIs

yashk2810 committed 29 days ago

9d4cc995

Introduce `manual_type: ManualAxisType` parameter on ShapedArray to track varying/unreduced/reduced and remove `vma: frozenset` parameter.

yashk2810 committed 29 days ago

ba3a7845

[hijax][boxes] fix nested jit qdd issue

mattjj committed 29 days ago

6b9d12af

Error out in lax.tile if a dimension is sharded and the rep is non-trivial (i.e. not 1).

yashk2810 committed 29 days ago

0f87a2fc

[pallas] rewrite pull_block_spec rule to work in terms of block_index transforms

levskaya committed 29 days ago

a628f21c

Merge pull request #35879 from gnecula:export_no_hlo_shardings

Google-ML-Automation committed 29 days ago

9a42fe6d

Guard tests that fail in libtpu==0.0.37, but succeed at nightly libtpu==0.0.38.dev20260315.

danielsuo committed 29 days ago

dcedfea2

Merge pull request #35974 from jakevdp:dct-complex

Google-ML-Automation committed 29 days ago

c7188d82

[pallas:triton] Fix atomic min/max lowering for unsigned integers and float types

draganmladjenovic committed 29 days ago

00d64246

Merge pull request #35610 from mattjj:remat-is-confusing

Google-ML-Automation committed 29 days ago

66575867

[hijax][remat3] enable dce optimization of remat3 vjp bwd whatever

mattjj committed 29 days ago

37fe2cb9

jax.scipy.fft: handle complex inputs to dct and friends

jakevdp committed 29 days ago

160ab443

Merge pull request #35967 from gnecula:export_doc2

Google-ML-Automation committed 29 days ago

297a27dd

Merge pull request #35972 from mattjj:primal-tangent-dtype-equality

Google-ML-Automation committed 29 days ago

93915eec

Merge pull request #35944 from ROCm:rocm-conftest-device

Google-ML-Automation committed 29 days ago

ca7e1f86

Introduce a new hypothesis_test_util module.

hawkinsp committed 29 days ago

0b871b0b

fix primal_tangent_dtype equality check

mattjj committed 29 days ago

71480ae5

Update docs/export/export.md

gnecula committed 29 days ago

Verified 4cb68bfc

[Mosaic] Prove divisibility of loop induction vars

naummo committed 29 days ago

3ad2cd96

[export] Clarified the jax.Export compatibility guarantees.

gnecula committed 29 days ago

f5b1e543

Disable GPU memory preallocation in ROCm pytest

gulsumgudukbay committed 29 days ago

d2d242ad

Merge pull request #35965 from gnecula:export_doc

Google-ML-Automation committed 29 days ago

bc9bbe70

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

gnecula committed 29 days ago

d48994d7

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

Google-ML-Automation committed 29 days ago

7c5f74c8

Newer Older