CI: 06/20/25 upstream sync#477

Open

rocm-repo-management-api-2[bot] wants to merge 2162 commits intomainfrom

ci-upstream-sync-220_1

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

Commits on Jun 6, 2025

Add fast-path for non-concrete Tracers in is_constant_dim to lower trace-time overhead from Exception creation

sharadmv
authored and
Google-ML-Automation
committed
Reverts 5c33588b30edbae51d5b63b0bd7cc8d9058d7ccb

Joshua Bambrick
authored and
Google-ML-Automation
committed
[Mosaic GPU] Extract the type-related logic out of `reinterpret_smem_ref`.

dimitar-asenov
authored and
Google-ML-Automation
committed
Port PartitionSpec to C++.

hawkinsp
authored and
Google-ML-Automation
committed
Fix segfault if None is passed to PartitionSpec.__eq__.

hawkinsp
authored and
Google-ML-Automation
committed
[Pallas][Mosaic GPU] Expose partitioned collective loads to copy_gmem_to_smem.

justinjfu
authored and
Google-ML-Automation
committed
Optimize jaxpr equation pretty-printing.

hawkinsp
authored and
Google-ML-Automation
committed
[JAX] Allow registering callbacks to be called when backends are cleared

hyeontaek
authored and
Google-ML-Automation
committed
PRNGKeyArray doesn't have a format field so assign the layout to be None when we are discharging the ref.

yashk2810
authored and
Google-ML-Automation
committed
Reverts c1bb095c5ce5b0286dc5052abf3b597b6f23cea5

mattjj
authored and
Google-ML-Automation
committed
Add committed property to MutableArray

yashk2810
authored and
Google-ML-Automation
committed
[Pallas TPU] Add custom_vjp_call lowering rule

sharadmv
authored and
Google-ML-Automation
committed
Small speedups to pretty-printing.

hawkinsp
authored and
Google-ML-Automation
committed
Set `in_sharding` to UNSPECIFIED if a mutableArray is uncommitted when a mutable array is closed over

yashk2810
authored and
Google-ML-Automation
committed

Commits on Jun 7, 2025

Commits on Jun 9, 2025

Commits on Jun 10, 2025

Commits on Jun 11, 2025

Commits on Jun 12, 2025

Commits on Jun 13, 2025

Use a frozenset for unconstrained_dims in sharding_constraint_p.

hawkinsp
authored and
Google-ML-Automation
committed
[jaxlib] Change Traceback to be a raw CPython class rather than a nanobind class.

hawkinsp
authored and
Google-ML-Automation
committed
Make the params of more jaxpr primitives hashable.

hawkinsp
authored and
Google-ML-Automation
committed
Remove unused internal optimization_barrier alias

Jake VanderPlas
authored and
Google-ML-Automation
committed
fix-forward for pallas tpu memory spaces test

sharadmv
authored and
Google-ML-Automation
committed
Update XLA dependency to use revision
Google-ML-Automation
committed
Move jax._src.callback to its own BUILD rule

Jake VanderPlas
authored and
Google-ML-Automation
committed
[Mosaic GPU] Convert all memrefs with transforms to unrealized casts and check them.

dimitar-asenov
authored and
Google-ML-Automation
committed
[Mosaic GPU] Add a Mosaic GPU op `with_transforms` for manually setting memref transforms.

dimitar-asenov
authored and
Google-ML-Automation
committed
[Mosaic GPU] Resolve different tile transforms using the largest common divisor.

dimitar-asenov
authored and
Google-ML-Automation
committed
[Mosaic GPU] Use warpgroup semantics for the ragged dot example kernel.

dimitar-asenov
authored and
Google-ML-Automation
committed
Disable `too_slow` in data.draw() for test_ndindexer

jakeharmon8
authored and
Google-ML-Automation
committed
[Mosaic GPU] Reconcile the swizzle of the a and b operands for wgmma in the Mosaic GPU dialect.

dimitar-asenov
authored and
Google-ML-Automation
committed
Add pjit_p to extend.core.primitives
j-towns
committed
Fix return type annotation for tree_util.tree_broadcast.

jburnim
authored and
Google-ML-Automation
committed
[Mosaic GPU] Parametrize the `test_subview` test.

dimitar-asenov
authored and
Google-ML-Automation
committed
Make params of several pallas primitives hashable.

hawkinsp
authored and
Google-ML-Automation
committed
Improve reshape not supported error message

yashk2810
authored and
Google-ML-Automation
committed
Internal refactor: move TPU lowering rules out of jax/_src/lax
jakevdp
committed
Make params of assert_consumed_value_p hashable.

hawkinsp
authored and
Google-ML-Automation
committed
Merge pull request #29456 from j-towns:pjit-p-extend-primitives
Google-ML-Automation
committed
Merge pull request #28931 from olupton:cublas-cudnn
Google-ML-Automation
committed
Merge pull request #29420 from jakevdp:tpu-linalg
Google-ML-Automation
committed
Make some remaining jaxpr equation params hashable.

hawkinsp
authored and
Google-ML-Automation
committed
Implemented cross-host memory transfer on GPU.

mwhittaker
authored and
Google-ML-Automation
committed
Load CUDA libraries up front with cdll.LoadLibrary().
hawkinsp
committed
Merge pull request #29462 from hawkinsp:nvrtc
Google-ML-Automation
committed
[Pallas] Add no_pipelining debugging option to emit_pipeline.

justinjfu
authored and
Google-ML-Automation
committed
Replace `with_partitions` and `with_unreduced` with `.update` on Partitions

yashk2810
authored and
Google-ML-Automation
committed
Remove `with_spec` from NamedSharding and replace with `.update`

yashk2810
authored and
Google-ML-Automation
committed

Commits on Jun 14, 2025

Commits on Jun 15, 2025

Update XLA dependency to use revision
Google-ML-Automation
committed

Commits on Jun 16, 2025

Update XLA dependency to use revision
Google-ML-Automation
committed
PR #28102: Add cudnn paged attention support in JAX cuDNN SDPA API

Cjkkkk
authored and
Google-ML-Automation
committed
Add `update_vma` and `update_weak_type` override on AbstractTMEMRef so that it return an instance of AbstractTMEMRef and not AbstractMemoryRef

yashk2810
authored and
Google-ML-Automation
committed
Add version guards to testAutoPgle

yashk2810
authored and
Google-ML-Automation
committed
Bump the libtpu check to 6/20

berkinilbeyi
authored and
Google-ML-Automation
committed
Removed unused `PyTreeDef::MakeFromNodeDataAndChildren` and its Python binding

superbobry
authored and
Google-ML-Automation
committed
[jax] Increase absolute test tolerance for lax_control_flow test

basioli-k
authored and
Google-ML-Automation
committed
Fix some more instances of unhashable jaxpr equation arguments.

hawkinsp
authored and
Google-ML-Automation
committed
[export] Add back-compat test for tridiagonal solve on GPU

gnecula
authored and
Google-ML-Automation
committed
Set heartbeat_timeout argument and flag.

mwhittaker
authored and
Google-ML-Automation
committed
Install SciPy from its source (head) to test against Python 3.14.0b1

kanglant
authored and
Google-ML-Automation
committed
[doc] add missing axis_types documentation
jakevdp
committed
Remove legacy CPU custom calls.

dfm
authored and
Google-ML-Automation
committed
Merge pull request #29497 from jakevdp:make-mesh-doc
Google-ML-Automation
committed
Fix a missing bounds check in traceback code.

hawkinsp
authored and
Google-ML-Automation
committed
Removed fixed suppressions
vfdev-5
authored
[Mosaic:TPU][NFC] Delete unused variable

tlongeri
authored and
Google-ML-Automation
committed
[JAX] Relax the return type of `colocated_python` decorator

hyeontaek
authored and
Google-ML-Automation
committed

Commits on Jun 17, 2025

Commits on Jun 18, 2025

Commits on Jun 19, 2025

Commits on Jun 20, 2025