Commits ROCm/jax - SemanticDiff

[Mosaic GPU][NFC] Use `dataclasses.fields` to gather operands for `CustomPrimitiveOp`.

allanrenucci committed 117 days ago

814818b3

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

Google-ML-Automation committed 117 days ago

265eacd5

Add core.standard_insert_vary before binding lt_to_p and also in searchsorted

yashk2810 committed 118 days ago

48335107

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0b874e4d99201cb0349483999cb9057c6c468d7f

Google-ML-Automation committed 118 days ago

8dd92e52

[Mosaic] Add `sign_extended` attribute to `tpu::UnpackSubelementsOp`.

WindQAQ committed 119 days ago

fbf4610f

Make eager shmap + check_vma=False + axis_index work with explicit mode.

yashk2810 committed 119 days ago

9855a2ee

Disallow nesting with jax.experimental.shard_map. This means to remove the `auto=` parameter from the experimental endpoint. Please use `jax.shard_map` if you want to nest shard_map calls.

yashk2810 committed 119 days ago

dfadef7e

Update XLA dependency to use revision http://github.com/openxla/xla/commit/23a18aa917e6ed7cedb8b9403f751dc5617a979c

Google-ML-Automation committed 119 days ago

bae1fbd5

Add a fancy transpose rule for `reshard_p` so that we can do the accum under the right context on the backward pass which is the source's context.

yashk2810 committed 119 days ago

c7557012

[JAX] Make Colocated Python object destruction use a sorted device list

hyeontaek committed 120 days ago

5cdc116c

Merge pull request #32535 from jakevdp:pallas-shard-map

Google-ML-Automation committed 120 days ago

16a8ca43

[IFRT] Migrate `Array::pjrt_layout()` callers to interpret `nullptr` as a default layout

hyeontaek committed 120 days ago

e64fe932

Update `pip` and `uv` default index with mirror artifacts repository path in Postsubmit JAX workflows.

ybaturina committed 120 days ago

653a8461

bump rules_ml_toolchain version

ermilovmaxim committed 120 days ago

164fcc15

Reverts e2a4c1a57a6942d4432fbb7ca9ddb0d9377cb49c

subhankarshah committed 120 days ago

d2391fcc

[pallas] use jax.shard_map in distributed test

jakevdp committed 120 days ago

c3924009

Remove jax.experimental.shard_map docstring

yashk2810 committed 120 days ago

6e85aa4e

Merge pull request #32527 from jax-ml:shmap

Google-ML-Automation committed 120 days ago

c15d1075

[bazel] add visibility hook for legacy `jax.experimental.shard_map`

Jake VanderPlas committed 120 days ago

39183b42

Use factory decorator pattern in shard_map tutorial

yashk2810 committed 120 days ago

58034a71

Merge pull request #32522 from jakevdp:shard-map-docs

Google-ML-Automation committed 120 days ago

0a429537

Remove docs for legacy jax.experimental.shard_map

jakevdp committed 120 days ago

e54a818a

Update `pip` and `uv` default index with mirror artifacts repository path in Presubmit JAX workflows.

ybaturina committed 120 days ago

f99c3e6a

Document that negative low/high padding is allowed for lax.pad.

hawkinsp committed 120 days ago

6c64f37a

Switch JAX to use an FFI call for cholesky decompositions rather than the Cholesky HLO on GPU.

hawkinsp committed 120 days ago

cf9a3528

[Mosaic GPU] Support splitting of minor dimension for collective copies when there is no swizzling.

allanrenucci committed 120 days ago

83d83044

Merge pull request #32499 from mattjj:issue32479

Google-ML-Automation committed 120 days ago

c4cfff9d

Merge pull request #32493 from mattjj:newstyle-hiprimitive2

Google-ML-Automation committed 120 days ago

05b4ac80

[Pallas:MGPU] Implement a basic reduce-scatter collective using multimem.ld_reduce

apaszke committed 120 days ago

97d2e204

[Mosaic GPU][NFC] Save 5 LOC.

allanrenucci committed 120 days ago

2fb239dd