ROCm/jax

Pull Requests Commits

Cherry-pick fixes for JAX 0.9.2 release from JAX 0.9.1 (#740)

magaonka-amd committed 14 hours ago

Verified dbc860fe

Prepare for JAX release 0.9.2

danielsuo committed 5 days ago

a659757d

Cherry-pick upstream XLA and Shardy fixes via Bazel patches.

danielsuo committed 5 days ago

e0bb88af

Guard tests that fail in libtpu==0.0.37, but succeed at nightly libtpu==0.0.38.dev20260315.

danielsuo committed 6 days ago

43195de1

Make sure to use `to_ct_spec()` when unsharding Zeros in shard_map_transpose. Also use `to_ct_aval` in _flatten_bwd in custom_vjp

yashk2810 committed 6 days ago

6ca53832

[Mosaic GPU] Add Pallas lowering for AsyncStoreSmemToTmem op.

golechwierowicz committed 8 days ago

ed8e9d55

[Mosaic GPU] Make tests involving global state order-independent.

bchetioui committed 8 days ago

a86c180e

[NFC] Reduce nesting with early return in `_core_map_lowering_rule`.

allanrenucci committed 8 days ago

7f265b45

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

Google-ML-Automation committed 8 days ago

e8253978

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

Google-ML-Automation committed 9 days ago

8304fec2

Merge pull request #34332 from jakevdp:searchsorted-hijax

Google-ML-Automation committed 9 days ago

13dcc922

Reduce the number of pytest workers for Linux aarch64 CPU builds.

hawkinsp committed 9 days ago

6463906c

Use a np.printoptions context manager rather than np.set_printoptions in test utility.

hawkinsp committed 9 days ago

5c4dd597

Reverts 3b4bee803de1a83b98cfe1aced47ab58f9bbc0c3

hawkinsp committed 9 days ago

7a68c4b1

jakevdp committed 9 days ago

5d7b5186

validate out_dtype

jakevdp committed 9 days ago

250ced1d

more fixes for sharded computation

jakevdp committed 9 days ago

aad9769e

Use scan rather than fori_loop

jakevdp committed 9 days ago

6ef7e8d8

[hijax] add hijax searchsorted primitive

jakevdp committed 9 days ago

2185cb1c

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

Google-ML-Automation committed 10 days ago

4ea536c2

Merge pull request #35696 from davisyoshida:recip-docstring

Google-ML-Automation committed 10 days ago

ea4b2e7e

davisyoshida committed 10 days ago

a767ad38

Fix reciprocal docstring

davisyoshida committed 10 days ago

f5676668

Automated Code Change

Google-ML-Automation committed 10 days ago

2d5a40fe

Make _get_owner skip sharding_constraints so that the xla_metadata is applied to the op instead of the sharding constraints

yashk2810 committed 10 days ago

a364e282

[Take 2] shard map without linear util

dougalm committed 10 days ago

830e342c

Perform dtype canonicalization during Primitive.bind.

hawkinsp committed 10 days ago

fae1e236

Merge pull request #35818 from jakevdp:lax-accuracy

Google-ML-Automation committed 10 days ago

9915f8c2

Merge pull request #35852 from ROCm:jax-rocm-version

Google-ML-Automation committed 10 days ago

964e2ebc

Merge pull request #35866 from jakevdp:take-current-trace

Google-ML-Automation committed 10 days ago

a8da2f48

Older