[Test do not merge] Testing build wheels #722

Closed

alekstheod wants to merge 7513 commits intomainfrom

This pull request is big! We're only showing the most recent 250 commits

Commits on Feb 20, 2026

Bump actions/checkout from 6.0.0 to 6.0.1
dependabot[bot]
authored

Commits on Feb 21, 2026

Propagate input sharding to output in full_like if it's fully replicated. Then even if input/output shapes don't match, it should be fine. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810
authored and
Google-ML-Automation
committed
[hijax] add is_vjp bit to LinearizeTrace

mattjj
and
yashk2810
committed
Merge pull request #35285 from mattjj:is-vjp-bit
Google-ML-Automation
committed
Update XLA dependency to use revision http://github.com/openxla/xla/commit/7505d7e818868a19f83fb0656444741ca114e444
Google-ML-Automation
committed
Canonicalize indexer.slice_sharding. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810
authored and
Google-ML-Automation
committed

Commits on Feb 22, 2026

Commits on Feb 23, 2026

Commits on Feb 24, 2026

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0db288651430120812e4208bad5bac269a8f6a54
Google-ML-Automation
committed
[Mosaic GPU] Add support for swizzle=16 in async SMEM->TMEM copies

apaszke
authored and
Google-ML-Automation
committed
Reverts 42c458cc318706bbbefbb45a752da6615a7123d3

apaszke
authored and
Google-ML-Automation
committed
More Pyrefly fixes
superbobry
committed
[Pallas:MGPU] Allow SMEM->TMEM copies with unswizzled (but tiled) refs

apaszke
authored and
Google-ML-Automation
committed
[Mosaic GPU] Relax N multiplicity restrictions for tcgen05 block-scaled MMA

apaszke
authored and
Google-ML-Automation
committed
[Mosaic GPU] Add support for fp4 tcgen05 MMAs that are both block-scaled and sparse

apaszke
authored and
Google-ML-Automation
committed
Merge pull request #35133 from mwhittaker:fault_tolerance_blog_tweaks
Google-ML-Automation
committed
[Mosaic TPU] Expose the latest version from the TPU MosaicSerdePass

apaszke
authored and
Google-ML-Automation
committed
Add a light, blocking Mosaic B200 presubmit.

belitskiy
authored and
Google-ML-Automation
committed
[Pallas:MGPU] Add support for writes to WGMMA accumulator refs

apaszke
authored and
Google-ML-Automation
committed
[Mosaic:GPU] Temporary fix of ASAN test.

PatriosTheGreat
authored and
Google-ML-Automation
committed
PR #35345: Bump actions/upload-artifact from 5.0.0 to 6.0.0

dependabot[bot]
authored and
Google-ML-Automation
committed
Merge pull request #35367 from superbobry:pyrefly
Google-ML-Automation
committed
[mosaic] Verify that `tpu.memref_slice` always produces tile-aligned slices

superbobry
authored and
Google-ML-Automation
committed
[hijax] mlir.lower_fun should have lo-jaxprs

mattjj
and
Jake Vanderplas
committed

Commits on Feb 25, 2026

Commits on Feb 26, 2026

Commits on Feb 27, 2026

Commits on Feb 28, 2026

Commits on Mar 1, 2026

Update XLA dependency to use revision http://github.com/openxla/xla/commit/15bc20b490170c25a4f4669d10573c6a601c0077
Google-ML-Automation
committed

Commits on Mar 2, 2026

Update XLA dependency to use revision http://github.com/openxla/xla/commit/ba5bbceae1ff6c0f03f0234ba6beadbcdae74635
Google-ML-Automation
committed
Add C addition to hopper matmul kernel.
Google-ML-Automation
committed
Merge remote-tracking branch 'origin/release/0.9.1' into postrelease/0.9.1
danielsuo
committed
Postrelease JAX v0.9.1.
danielsuo
committed
[jaxlib] :normalize_stubs now also fixes improper ``get`` classmethods
superbobry
committed
[mgpu] Minor cleanup in `try_cluster_cancel`.

chr1sj0nes
authored and
Google-ML-Automation
committed
[Mosaic GPU][NFC] Remove duplicate layout inference rule for `vector.BroadcastOp`.

allanrenucci
authored and
Google-ML-Automation
committed
Merge pull request #35502 from superbobry:stubgen
Google-ML-Automation
committed
[Mosaic GPU][NFC] Use more precise type for `op` in `_optimization_barrier_constraint_system`.

allanrenucci
authored and
Google-ML-Automation
committed
[pallas:sc] Switched to using `pltpu.emit_pipeline` in the lowering

superbobry
authored and
Google-ML-Automation
committed
[Mosaic] Add assembly format to tpu.delay_op

vsytch
authored and
Google-ML-Automation
committed
[Mosaic GPU] Fix for the latest LLVM version

apaszke
authored and
Google-ML-Automation
committed
Add serialization support for Mosaic GPU kernels.

khasanovaa
authored and
Google-ML-Automation
committed
Regenerate type stubs for nanobind v2.12.0.

hawkinsp
authored and
Google-ML-Automation
committed
[Mosaic GPU][NFC] Remove dead `constraints.BroadcastInDim`.

allanrenucci
authored and
Google-ML-Automation
committed
Merge pull request #35501 from jax-ml:postrelease/0.9.1
Google-ML-Automation
committed
Reverts 41ae1eb6ea50739feb91589705aa078b1495d40f

hawkinsp
authored and
Google-ML-Automation
committed
Merge pull request #34702 from ROCm:add-run-bazel-test-rocm-rbe
Google-ML-Automation
committed
Merge pull request #35417 from joshuapjacob:patch-1
Google-ML-Automation
committed
[Mosaic GPU] Add support for untiled CP_ASYNC

apaszke
authored and
Google-ML-Automation
committed
Merge pull request #35453 from jakevdp:pyrefly-ad
Google-ML-Automation
committed
Remove nvidia_wheel_versions

charleshofer
authored and
mminutoli
committed
Make jaxlib targets visible

charleshofer
authored and
mminutoli
committed
hipblas typedef fix

charleshofer
authored and
mminutoli
committed
No GPU fail

charleshofer
authored and
mminutoli
committed
Wrap HIP inline functions in anonymous namespaces in vendor.h
mminutoli
committed
SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd
authored and
mminutoli
committed
Add shared utility function get_rocm_version to test_util.py

charleshofer
authored and
mminutoli
committed
Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin
authored and
mminutoli
committed
Fix v_pages quantization and adjust test params for ROCm compatibilit… (#560)

phambinhfin
authored and
mminutoli
committed
Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563)

Arech8
authored and
mminutoli
committed
Add skip of test_is_finite() on Cuda (#565)

Arech8
authored and
mminutoli
committed
Add rocm test requirements file (#570)

AratiGanesh
authored and
mminutoli
committed
Let the unit tests use build.py for setting up Bazel commands for unit tests (#582)

charleshofer
authored and
mminutoli
committed
adding abort logic to rocm/jax (#590)

gulsumgudukbay
authored and
mminutoli
committed
Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (#597)

phambinhfin
authored and
mminutoli
committed
Fix shared memory limit check for ROCm in test_dot (#596)

phambinhfin
authored and
mminutoli
committed
Fix Numpy signatures test (#598)

authored and
mminutoli
committed
fix merge arts

Ruturaj4
authored and
mminutoli
committed
Enable RngShardingTests (#644)

gulsumgudukbay
authored and
mminutoli
committed
Enable test_variadic_reduce_window on ROCm (#647)
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Enable testMultivariateNormalSingularCovariance on ROCm (#666)

AratiGanesh
authored and
mminutoli
committed
Skip test_tridiagonal_solve on ROCm due to hipSPARSE numerical errors (#668)

AratiGanesh
authored and
mminutoli
committed
Update Skip Reason Outputs (#663)

gulsumgudukbay
authored and
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Skip testCudaArrayInterfaceOnNonCudaFails on ROCm platform (#677)

magaonka-amd
authored and
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
mminutoli
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
mminutoli
committed
Add ROCm encoding for test_struct_encoding_determinism (#683)

AratiGanesh
authored and
mminutoli
committed
Remove 'mean' from unsupported params for jnp.var (#689)

magaonka-amd
authored and
mminutoli
committed
Implement approx_tanh for ROCm using OCML tanh function (#691)

magaonka-amd
authored and
mminutoli
committed
Skipping testEighTinyNorm due to hipSolver issues (#697)

AratiGanesh
authored and
mminutoli
committed
Abort detection CI workflow (#688)

gulsumgudukbay
authored and
mminutoli
committed
Abort-Detection: Fix halt-for-connection input (#712)

gulsumgudukbay
authored and
mminutoli
committed