ROCm/jax

Pull Requests Commits

aiter mha kernels (ASM+CK) integration

zahiqbal committed 2 hours ago

a6041a3b

Add support for parallelizing lapack invocations on the batch dimension.

hawkinsp committed 8 hours ago

bfd9e1fb

[Mosaic TPU] Allow inferring divisibility through divisions

apaszke committed 8 hours ago

3ba0e1f2

[Pallas:MGPU] Make sure that the warp axis returns indices in [0, 3)

apaszke committed 8 hours ago

5bd8a46e

Merge pull request #36559 from superbobry:main

Google-ML-Automation committed 9 hours ago

e0696314

Upgraded to Pyrefly 0.60.0

superbobry committed 9 hours ago

34bf5455

[Mosaic GPU] Enable layout inference for collective scaled MMA.

allanrenucci committed 11 hours ago

03cf9812

[mGPU] Use CPU target config, that is passed from XLA side via FFI.

khasanovaa committed 11 hours ago

8708d738

add a test for a simple elementwise unary kernel

ImanHosseini committed 11 hours ago

6e73777c

[Mosaic GPU] Add an `AnyOf` constraint to layout inference.

allanrenucci committed 12 hours ago

7141987a

Merge pull request #36549 from jakevdp:pyrefly-059

Google-ML-Automation committed 12 hours ago

cbc13839

[pallas:mgpu] Allow `commit_smem_to_gmem_group` in a single warp context.

chr1sj0nes committed 13 hours ago

9bf6adf3

Remove the tools/build_defs.bzl

gnecula committed 13 hours ago

254515d8

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

Google-ML-Automation committed 13 hours ago

8552ae99

[Mosaic GPU][NFC] Add `orders_tensor_core` to DialectBarrierRef.

allanrenucci committed 14 hours ago

beed3313

Merge pull request #36525 from gnecula:consts_2

Google-ML-Automation committed 16 hours ago

8c1da538

Make sure vary_unreduced_cast and reduced_vary_cast order the axes wrt the context mesh. Also code-golf some code in shard_map

yashk2810 committed 21 hours ago

6f55c374

Merge pull request #36551 from jakevdp:eig-autodiff-docs

Google-ML-Automation committed 23 hours ago

a6285551

[pmap] Create new ifrt array when rewrapping global arrays as host local in multihost pmap.

danielsuo committed 23 hours ago

6bad8fcf

Merge pull request #36535 from jakevdp:numpy-dtype-dep

Google-ML-Automation committed 23 hours ago

1b046188

[doc] update docs for eig() autodiff

jakevdp committed 23 hours ago

18623265

[pyrefly] update pyrefly to v0.59.1 & fix/silence errors

jakevdp committed 1 day ago

f1935ff1

Support TransformedRefs in higher-order primitives by tree flattening them.

rdyro committed 1 day ago

60a8d25b

Tighten the type signature of jax.numpy.where

Jake VanderPlas committed 1 day ago

74ab31ee

Reverts aaec4d1ef9568ac5bea7d48497fde996aafc8402

hawkinsp committed 1 day ago

e5ca8d47

Experimental pallas-triton lowering of ragged_dot_general on GPU.

rdyro committed 1 day ago

9a9fc023

Merge pull request #36537 from jakevdp:version-support

Google-ML-Automation committed 1 day ago

4d108a60

Update NumPy/SciPy version support

jakevdp committed 1 day ago

08aedc7c

[maint] add named deprecation for array dtype conversion

jakevdp committed 1 day ago

e2f64152

[JAX] Support non-PjRt-compatible Arrays in JAX heap profile API

hyeontaek committed 1 day ago

8c197055

Older