ROCm/jax

Pull Requests Commits

Add support for parallelizing lapack invocations on the batch dimension.

hawkinsp committed 4 days ago

bfd9e1fb

[Mosaic TPU] Allow inferring divisibility through divisions

apaszke committed 4 days ago

3ba0e1f2

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

apaszke committed 4 days ago

5bd8a46e

Merge pull request #36559 from superbobry:main

Google-ML-Automation committed 4 days ago

e0696314

Upgraded to Pyrefly 0.60.0

superbobry committed 4 days ago

34bf5455

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

allanrenucci committed 4 days ago

03cf9812

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

khasanovaa committed 4 days ago

8708d738

add a test for a simple elementwise unary kernel

ImanHosseini committed 4 days ago

6e73777c

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

allanrenucci committed 4 days ago

7141987a

Merge pull request #36549 from jakevdp:pyrefly-059

Google-ML-Automation committed 4 days ago

cbc13839

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

chr1sj0nes committed 4 days ago

9bf6adf3

Remove the tools/build_defs.bzl

gnecula committed 4 days ago

254515d8

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

Google-ML-Automation committed 4 days ago

8552ae99

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

allanrenucci committed 4 days ago

beed3313

Merge pull request #36525 from gnecula:consts_2

Google-ML-Automation committed 5 days 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 5 days ago

6f55c374

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

Google-ML-Automation committed 5 days ago

a6285551

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

danielsuo committed 5 days ago

6bad8fcf

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

Google-ML-Automation committed 5 days ago

1b046188

[doc] update docs for eig() autodiff

jakevdp committed 5 days ago

18623265

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

jakevdp committed 5 days ago

f1935ff1

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

rdyro committed 5 days ago

60a8d25b

Tighten the type signature of jax.numpy.where

Jake VanderPlas committed 5 days ago

74ab31ee

Reverts aaec4d1ef9568ac5bea7d48497fde996aafc8402

hawkinsp committed 5 days ago

e5ca8d47

Experimental pallas-triton lowering of ragged_dot_general on GPU.

rdyro committed 5 days ago

9a9fc023

Merge pull request #36537 from jakevdp:version-support

Google-ML-Automation committed 5 days ago

4d108a60

Update NumPy/SciPy version support

jakevdp committed 5 days ago

08aedc7c

[maint] add named deprecation for array dtype conversion

jakevdp committed 5 days ago

e2f64152

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

hyeontaek committed 5 days ago

8c197055

Remove the jax_to_tf build rule and the tfcompile folder

gnecula committed 5 days ago

33dae7ee

Older