ROCm/jax

Pull Requests Commits

Move ROCm GPU detection to hardware_utils and add multi-GPU shm warning

alekstheod committed 21 days ago

84b90c8e

[Mosaic GPU] Add TMEM aliasing support for WG semantics.

dimitar-asenov committed 21 days ago

521f7951

Add regression test of https://github.com/jax-ml/jax/pull/35787

GleasonK committed 21 days ago

a231919a

Replace `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`. Also expose `jax.sharding.ManualAxisType` as a public API.

yashk2810 committed 21 days ago

2ed3944c

[Mosaic GPU] Improve tcgen05 preconditions error messaging.

golechwierowicz committed 21 days ago

93772b38

[NFC] Streamline `copy_gmem_to_smem_p` lowering registration.

allanrenucci committed 21 days ago

f6fd0352

[NFC] Streamline lowering registration for warp primitive semantics.

allanrenucci committed 21 days ago

96531f06

[Mosaic GPU] Add Pallas WG lowering for async_copy_sparse_metadata_to_tmem.

golechwierowicz committed 21 days ago

e4d133a8

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

Google-ML-Automation committed 21 days ago

464b047e

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

ZixuanJiang committed 22 days ago

3558a4af

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

yashk2810 committed 22 days ago

9d4cc995

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

yashk2810 committed 22 days ago

ba3a7845

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

yashk2810 committed 22 days ago

0f87a2fc

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

levskaya committed 22 days ago

a628f21c

Merge pull request #35879 from gnecula:export_no_hlo_shardings

Google-ML-Automation committed 22 days ago

9a42fe6d

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

danielsuo committed 22 days ago

dcedfea2

Merge pull request #35974 from jakevdp:dct-complex

Google-ML-Automation committed 22 days ago

c7188d82

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

Google-ML-Automation committed 22 days ago

66575867

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

mattjj committed 22 days ago

37fe2cb9

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

jakevdp committed 22 days ago

160ab443

Merge pull request #35967 from gnecula:export_doc2

Google-ML-Automation committed 22 days ago

297a27dd

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

Google-ML-Automation committed 22 days ago

93915eec

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

Google-ML-Automation committed 22 days ago

ca7e1f86

Introduce a new hypothesis_test_util module.

hawkinsp committed 22 days ago

0b871b0b

fix primal_tangent_dtype equality check

mattjj committed 22 days ago

71480ae5

Update docs/export/export.md

gnecula committed 22 days ago

Verified 4cb68bfc

[Mosaic] Prove divisibility of loop induction vars

naummo committed 22 days ago

3ad2cd96

[export] Clarified the jax.Export compatibility guarantees.

gnecula committed 22 days ago

f5b1e543

Disable GPU memory preallocation in ROCm pytest

gulsumgudukbay committed 22 days ago

d2d242ad

Merge pull request #35965 from gnecula:export_doc

Google-ML-Automation committed 22 days ago

bc9bbe70

Older