ROCm/jax

Pull Requests Commits

Delete stale xla_gpu_cublaslt_algorithm3.patch

srinivamd committed 5 days ago

467dca54

Remove stale xla_gpu_cublaslt_algorithm3.patch from workspace.bzl

srinivamd committed 5 days ago

ecb65827

Bump pinned jaxlib in requirements_lock files from 0.9.0 to 0.9.2 (#769)

magaonka-amd committed 15 days ago

Verified 877480e3

[ROCm] Bump XLA pin to ROCm/xla rocm-jaxlib-v0.9.2 HEAD (rocprofiler-sdk v3 + roctracer v1 build fix) (#766)

magaonka-amd committed 20 days ago

Verified aae3e281

[ROCm] Cherry-pick: Fix inter-block write race in pallas test_non_range_while_loop (#764)

magaonka-amd committed 21 days ago

Verified c7a70f49

[ROCm] Cherry-pick: Split ROCm pytest into single- and multi-accelerator passes (#760)

magaonka-amd committed 21 days ago

Verified bde64f36

[ROCm] Cherry-pick: Override HIP_VISIBLE_DEVICES per ROCm xdist worker (#761)

magaonka-amd committed 21 days ago

Verified c0ddc7f3

[ROCm] Cherry-pick: Fix env var leak in xla_bridge_test register_plugin tests (#762)

magaonka-amd committed 21 days ago

Verified dc55cc48

Cherry-pick: Fix visibility in BUILD files (#757) (#759)

magaonka-amd committed 21 days ago

Verified 3992836e

[ROCm] Cherry-pick: Add ROCm xdist device pinning in conftest.py (#744)

magaonka-amd committed 48 days ago

Verified c7fb7066

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

magaonka-amd committed 58 days ago

Verified dbc860fe

Prepare for JAX release 0.9.2

danielsuo committed 63 days ago

a659757d

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

danielsuo committed 63 days ago

e0bb88af

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

danielsuo committed 64 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 64 days ago

6ca53832

[Mosaic GPU] Add Pallas lowering for AsyncStoreSmemToTmem op.

golechwierowicz committed 65 days ago

ed8e9d55

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

bchetioui committed 65 days ago

a86c180e

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

allanrenucci committed 65 days ago

7f265b45

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

Google-ML-Automation committed 65 days ago

e8253978

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

Google-ML-Automation committed 66 days ago

8304fec2

Merge pull request #34332 from jakevdp:searchsorted-hijax

Google-ML-Automation committed 67 days ago

13dcc922

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

hawkinsp committed 67 days ago

6463906c

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

hawkinsp committed 67 days ago

5c4dd597

Reverts 3b4bee803de1a83b98cfe1aced47ab58f9bbc0c3

hawkinsp committed 67 days ago

7a68c4b1

jakevdp committed 67 days ago

5d7b5186

validate out_dtype

jakevdp committed 67 days ago

250ced1d

more fixes for sharded computation

jakevdp committed 67 days ago

aad9769e

Use scan rather than fori_loop

jakevdp committed 67 days ago

6ef7e8d8

[hijax] add hijax searchsorted primitive

jakevdp committed 67 days ago

2185cb1c

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

Google-ML-Automation committed 67 days ago

4ea536c2

Older