ROCm/jax

Pull Requests Commits

code review updates

zahiqbal committed 145 days ago

ab93aac3

corrected mha asm_v3 vs ck kernel calls

zahiqbal committed 145 days ago

b616cf94

upgraded aiter libs

zahiqbal committed 146 days ago

121df61e

handling aiter-mha libs in wheel at rocm/jax fork

zahiqbal committed 146 days ago

f461fcf3

code cleanup and formatting after code review

zahiqbal committed 150 days ago

720431d3

aiter mha kernels (ASM+CK) integration

zahiqbal committed 160 days ago

a6041a3b

Add support for parallelizing lapack invocations on the batch dimension.

hawkinsp committed 161 days ago

bfd9e1fb

[Mosaic TPU] Allow inferring divisibility through divisions

apaszke committed 161 days ago

3ba0e1f2

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

apaszke committed 161 days ago

5bd8a46e

Merge pull request #36559 from superbobry:main

Google-ML-Automation committed 161 days ago

e0696314

Upgraded to Pyrefly 0.60.0

superbobry committed 161 days ago

34bf5455

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

allanrenucci committed 161 days ago

03cf9812

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

khasanovaa committed 161 days ago

8708d738

add a test for a simple elementwise unary kernel

ImanHosseini committed 161 days ago

6e73777c

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

allanrenucci committed 161 days ago

7141987a

Merge pull request #36549 from jakevdp:pyrefly-059

Google-ML-Automation committed 161 days ago

cbc13839

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

chr1sj0nes committed 161 days ago

9bf6adf3

Remove the tools/build_defs.bzl

gnecula committed 161 days ago

254515d8

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

Google-ML-Automation committed 161 days ago

8552ae99

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

allanrenucci committed 161 days ago

beed3313

Merge pull request #36525 from gnecula:consts_2

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

6f55c374

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

Google-ML-Automation committed 161 days ago

a6285551

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

danielsuo committed 161 days ago

6bad8fcf

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

Google-ML-Automation committed 161 days ago

1b046188

[doc] update docs for eig() autodiff

jakevdp committed 161 days ago

18623265

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

jakevdp committed 161 days ago

f1935ff1

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

rdyro committed 161 days ago

60a8d25b

Tighten the type signature of jax.numpy.where

Jake VanderPlas committed 161 days ago

74ab31ee

Reverts aaec4d1ef9568ac5bea7d48497fde996aafc8402

hawkinsp committed 161 days ago

e5ca8d47

Older