ROCm/jax

Pull Requests Commits

aiter mha kernels (ASM+CK) integration (#747)

zahiqbal committed 6 days ago

272222bb

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

Google-ML-Automation committed 10 days ago

35deac3e

Enter into `eqn.ctx` in checkify's jaxpr interpreter so that the correct context mesh is used.

yashk2810 committed 11 days ago

fa1538d2

Add a comment to .pyi files that they are automatically generated.

hawkinsp committed 11 days ago

ee30a5bc

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

Google-ML-Automation committed 11 days ago

b09aa725

[Pallas] Enable state effects within mpmd_map for passed in Refs.

sharadmv committed 11 days ago

1fcf9154

Merge pull request #37113 from mattjj:ref-vmap

Google-ML-Automation committed 11 days ago

2ed28ab6

[Mosaic TPU] Support reshapes like [src_major_dims, A, B] -> [dst_major_dims, B, A], where B is a multiple of A, and B <= 128.

yueshengys committed 11 days ago

bf7e063c

Merge pull request #37125 from vfdev-5:jrand-api--dirichlet-out_sharding

Google-ML-Automation committed 11 days ago

a81d241f

[NFC] [Pallas TPU] Refactor Pallas Mosaic lowering to support unpipelined kernels via a simpler lowering path

sharadmv committed 11 days ago

65cd8839

Merge pull request #37153 from vfdev-5:jrand-api--categ-out_sharding

Google-ML-Automation committed 11 days ago

e08ca5e0

Added out_sharding arg to jax.random.dirichlet sampler

vfdev-5 committed 12 days ago

85816b87

Added out_sharding arg to jax.random.categorical sampler

vfdev-5 committed 12 days ago

ed25c147

Merge pull request #37157 from vfdev-5:jrand-api--logistic-out_sharding

Google-ML-Automation committed 12 days ago

755b7cb9

mattjj committed 12 days ago

ab430e31

Merge pull request #37158 from vfdev-5:jrand-api--pareto-out_sharding

Google-ML-Automation committed 12 days ago

9e952e6b

[Pallas MGPU] Add warp-level support to semaphore_signal

Rifur13 committed 12 days ago

f24c3070

Reland _convert_element_type fast path with two changes:

hawkinsp committed 12 days ago

b19a213c

Merge pull request #37152 from vfdev-5:jrand-api--laplace-out_sharding

Google-ML-Automation committed 12 days ago

543c93c5

Merge pull request #37154 from aybchan:aybchan/update-mesh-context

Google-ML-Automation committed 12 days ago

7fc555cd

Merge pull request #37126 from vfdev-5:jrand-api--exp-out_sharding

Google-ML-Automation committed 12 days ago

ade2492b

Reverts 55f2f21fb3f2576e375884ab7f33d1d01e65dc29

Google-ML-Automation committed 12 days ago

0d1361ff

[gh] Zip test artifacts from github actions workflow run jobs.

danielsuo committed 12 days ago

fd5a0b8e

Removed more unused Pyrefly ignores

superbobry committed 12 days ago

a274ce66

Merge pull request #37103 from vfdev-5:jrand-api--rayleigh-out_sharding

Google-ML-Automation committed 12 days ago

f6e64da2

[Pallas:MGPU] Fix lowering for `layout_cast` of scalars under WG semantics.

allanrenucci committed 12 days ago

504d3d96

Added out_sharding to jax.random.pareto sampler

vfdev-5 committed 12 days ago

42221be6

Added out_sharding arg to jax.random.logistic sampler

vfdev-5 committed 12 days ago

6958e799

Update mesh context manager

aybchan committed 12 days ago

2e74935c

Added out_sharding arg to jax.random.laplace sampler

vfdev-5 committed 12 days ago

5f6bb61c

Older