ROCm/jax

Pull Requests Commits

Resolve undefined behavior in bitshift unit test

mminutoli committed 183 days ago

49497750

Remove usage of `.tar` files in CI jobs due to fast extraction of `tar.xz` files done by hermetic `xz` tool.

ybaturina committed 183 days ago

c25e095f

Add a docstring for jax.set_mesh

yashk2810 committed 183 days ago

0dfd3552

Use dtype=int instead of explicit dtype=jnp.int64 in transposed ragged dot MGPU kernel, so it works for both x64 and x32 configs.

Google-ML-Automation committed 183 days ago

7c6c2bb6

Allow reduced on the fwd pass to jnp.sin even though unreduced is not allowed as an input to sin since sin is a non-linear op.

yashk2810 committed 183 days ago

7448556f

Add instructions for updating jax wheels content.

ybaturina committed 183 days ago

9265a785

Relax `TpuInfo.is_matmul_supported` for TPU7x (Ironwood) to match flexibility already built in HW.

Google-ML-Automation committed 183 days ago

e150368a

[Mosaic] Relax the constraint of doubling contracting size.

WindQAQ committed 183 days ago

f4fddc20

Reverts 46fb39e4b5aec751485a019385b7a5753666f6f3

Google-ML-Automation committed 183 days ago

aae42a82

[Mosaic GPU] Add a constraint to ensure that gmem indices in async_load/store are aligned with the tile size.

dimitar-asenov committed 183 days ago

32138c0c

Ignore pytype errors produced with --use-functools-partial-overlay

superbobry committed 183 days ago

78fad69d

[Mosaic GPU] Adjust swizzle inference for `tcgen05_mma` with M=64.

allanrenucci committed 184 days ago

0a3f1042

[Pallas:MGPU][NFC] Style nit.

allanrenucci committed 184 days ago

ed53d7f6

Merge pull request #32910 from gnecula:repro_hooks

Google-ML-Automation committed 184 days ago

01e5d089

Add hooks for the experimental JAX repro module.

gnecula committed 184 days ago

dd432452

[Pallas:MGPU] Add support for integer WGMMA under WG semantic.

allanrenucci committed 184 days ago

3918e675

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

Google-ML-Automation committed 184 days ago

72450552

Add an optimization to `PreCanonicalizationOptimizationPass` that removes conversions from lower precision to higher precision of the intermediate computations can be done losslessly in low precision.

Google-ML-Automation committed 184 days ago

46fb39e4

[Pallas TPU] Add `pack_elementwise` and `unpack_elementwise` primitives.

megrez-yliu committed 184 days ago

445d75c1

[Mosaic:TPU] Bitcast vreg instead of mask in SelectWithBounds

tlongeri committed 184 days ago

2138cd08

Internal change

yashk2810 committed 184 days ago

dcd51793

Add support for HiType in Pallas Fuser API

Google-ML-Automation committed 184 days ago

133b1618

Add test for mgpu ragged dot with transpose on RHS

Google-ML-Automation committed 184 days ago

a3c9f67b

Transposed ragged dot kernel on Hopper GPU

Google-ML-Automation committed 184 days ago

eb2d56b9

Merge pull request #33050 from mattjj:remat-refs

Google-ML-Automation committed 184 days ago

e9421c6c

Merge pull request #33124 from jakevdp:all-to-all-impl

Google-ML-Automation committed 184 days ago

7c248873

[jax] Use explicit FFI type registration in GPU plugins

ezhulenev committed 184 days ago

f4844e6d

(ragged) all_to_all: more informative error in eager mode

jakevdp committed 184 days ago

3800ed31

[Mosaic GPU] Make an error message more precise.

dimitar-asenov committed 184 days ago

aba5e9f6

Merge pull request #33116 from jakevdp:top-k-axis

Google-ML-Automation committed 184 days ago

91dd6a7b

Older