ROCm/jax

Pull Requests Commits

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 159 days ago

7d927561

Merge pull request #38062 from hawkinsp:docs3

Google-ML-Automation committed 159 days ago

a5766401

Merge pull request #37977 from ROCm:use_hermetic_sysroot_for_jax_ut

Google-ML-Automation committed 159 days ago

96cd401e

Merge pull request #38010 from olupton:pbroadcast

Google-ML-Automation committed 159 days ago

54c30905

Merge pull request #37320 from samanklesaria:geometric_out_sharding

Google-ML-Automation committed 159 days ago

f941b14e

Make warnings about security on the compilation cache warnings.

hawkinsp committed 159 days ago

d5ed6d08

[pallas:mgpu] Do not use `pl.pallas_call` in :gpu_pallas_distributed_test

superbobry committed 159 days ago

f180cf4d

[pallas:sc] Do not use `mlir.lower_per_platform` for `mpmd_map`

superbobry committed 159 days ago

cf248ec2

[Mosaic GPU] Simulate commuting transforms in WG semantics.

golechwierowicz committed 159 days ago

f7fb9d27

Merge pull request #37987 from jakevdp:random-docs

Google-ML-Automation committed 159 days ago

416cb209

[NFC] Derive `shape` from `tmem_layout` in `test_load_store_tmem_scales_layout`

allanrenucci committed 159 days ago

0922eeab

Merge pull request #37254 from samanklesaria:fix_triangular

Google-ML-Automation committed 159 days ago

f9f06d42

[pallas:sc] Removed the lowering rule for `jaxpr_call`

superbobry committed 159 days ago

6fb39be2

Remove setting `Manual` HloSharding directly as it would be anyways ignored during Shardy import. If required, JAX NamedSharding should be set instead which would be correctly lowered to corresponding sdy shardings.

KanishAnand committed 159 days ago

096a1402

[Mosaic:GPU] Use multimem_parameters attribute to signal that parameter is allocated within the collective memory space.

PatriosTheGreat committed 159 days ago

17b80742

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

Google-ML-Automation committed 159 days ago

0fd5af23

[JAX] Enable CHLO lowering for cumulative sums on GPU again.

chsigg committed 159 days ago

8f711b48

Make sure scan transposes correctly when carry is reduced.

yashk2810 committed 160 days ago

2a576d79

[Mosaic TPU] Support static lane and major dims shift of `pltpu.roll` with zero stride for packed types.

yueshengys committed 160 days ago

a1ff1d64

Error out if `*pspec` or `pspec[...]` is used when unreduced/reduced is non-empty

yashk2810 committed 160 days ago

d6c717d9

Merge pull request #38043 from jax-ml:mattjj-patch-1

Google-ML-Automation committed 160 days ago

756523e0

Fix formatting of JAX Ref documentation

mattjj committed 160 days ago

Verified cddcebdc

Add out_sharding to segment_{sum,max,min,prod}. Only allow unreduced out_sharding to `segment_sum`.

yashk2810 committed 160 days ago

ba968d94

Merge pull request #38039 from jakevdp:fix-cumsum

Google-ML-Automation committed 160 days ago

6f201d8a

Fix formatting of JAX Ref documentation

mattjj committed 160 days ago

Verified a5c1c30d

[Pallas] Improve error handling in TPU Interpret Mode to avoid hangs.

jburnim committed 160 days ago

3aa5aa75

[bug] fix jnp.cumsum with dtype=bool

jakevdp committed 160 days ago

9d7ed1b0

[Mosaic GPU] Add extra layout inference test for `AsyncLoadTmemOp`.

allanrenucci committed 160 days ago

0198884c

[Mosaic GPU] Support load from TMEM (with scales layout) to register.

allanrenucci committed 160 days ago

7cc2f41b

Add pbroadcast to docs.

olupton committed 160 days ago

eab94af7

Older