Add iree_metal to platforms with buffer donation support#2

Open

robtaylor wants to merge 380 commits intomainfrom

add-iree-metal-donation-support

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

Commits on Dec 8, 2025

Change PjRt to use new copy of coordination service.

mwhittaker
authored and
Google-ML-Automation
committed

Commits on Dec 9, 2025

[Pallas SC] Allow semaphores to be returned by SCS kernels

sharadmv
authored and
Google-ML-Automation
committed
Deprecate `with mesh:` context manager. Use `with jax.set_mesh(mesh):` instead

yashk2810
authored and
Google-ML-Automation
committed
Add a jax network transfer benchmark script
Google-ML-Automation
committed
Add ResultHandler.wrap which allows us to avoid exploding PRNG keys.

pschuh
authored and
Google-ML-Automation
committed
[Pallas] Device Id dict to mesh fastpath for power of twos
Google-ML-Automation
committed
Update XLA dependency to use revision http://github.com/openxla/xla/commit/91b3f740b75d1d932a12fb0886338f84f856a453
Google-ML-Automation
committed
[Mosaic GPU] Remove `{Least,Most}ReplicatedExpression` constructors.

bchetioui
authored and
Google-ML-Automation
committed
[Mosaic GPU] Support `FragmentedArray.broadcast_in_dim` for splat to other layouts.

allanrenucci
authored and
Google-ML-Automation
committed
[maint] clean up jnp.arange implementation
jakevdp
committed
Update generated stubs.

cantonios
authored and
Google-ML-Automation
committed
Fix back compat test to ignore warnings

yashk2810
authored and
Google-ML-Automation
committed
[pallas:mosaic_gpu] Slightly tweaked the error messages in a few places

superbobry
authored and
Google-ML-Automation
committed
[export] Add backwards compatibility tests for serialized exports
gnecula
committed
Use ifrt::AttributeMap::Get instead of directly accessing map

krishnaharidasan
authored and
Google-ML-Automation
committed
Merge pull request #33758 from gnecula:export_back_compat
Google-ML-Automation
committed
Mention pmap is in maintenance mode and point to shard_map and the migration guide

authored and
Google-ML-Automation
committed
[mosaic] Added documentation and a few useful methods to tpu.tiled attribute

superbobry
authored and
Google-ML-Automation
committed
Remove multiprocess tests from TPU presubmit due to latency.

emilyfertig
authored and
Google-ML-Automation
committed
Adds docs on multi-controller JAX fault tolerance.
mwhittaker
committed
[test] fix signatures test for NumPy nightly
jakevdp
committed
Merge pull request #33819 from jakevdp:arange-followup
Google-ML-Automation
committed
[test] remove pre-NumPy 2.0 API skips
jakevdp
committed
Merge pull request #32218 from mwhittaker:fault_tolerance_docs
Google-ML-Automation
committed
Merge pull request #33834 from jakevdp:numpy-nightly-sig
Google-ML-Automation
committed
Remove dynamic shapes. Dead weight at this point.

dougalm
authored and
mattjj
committed
jnp.arange: deprecate passing complex arguments
jakevdp
committed
Merge pull request #33836 from jakevdp:old-np-apis
Google-ML-Automation
committed
Merge pull request #33565 from jax-ml:remove-dynamic-shapes
Google-ML-Automation
committed
Merge pull request #33837 from jakevdp:arange-complex-dep
Google-ML-Automation
committed
[Pallas TPU] Allow closed over scalars in core_map code

sharadmv
authored and
Google-ML-Automation
committed

Commits on Dec 10, 2025

Commits on Dec 11, 2025

Commits on Dec 12, 2025

Commits on Dec 13, 2025

Commits on Dec 14, 2025

Commits on Dec 15, 2025

Commits on Dec 16, 2025

Commits on Dec 17, 2025

Commits on Dec 18, 2025

Commits on Dec 19, 2025

anselm refs

mattjj
authored and
Google-ML-Automation
committed
[JAX] Refresh a custom layout if a buffer is copied across clients or memories

hyeontaek
authored and
Google-ML-Automation
committed
Update the requirements' lock files post the JAX 0.8.2 release.

belitskiy
authored and
Google-ML-Automation
committed
Fix multi_broadcast_in_dim which was doing replicated -> unreduced casts which was unnecessary.

yashk2810
authored and
Google-ML-Automation
committed
anselm refs

mattjj
authored and
Google-ML-Automation
committed
Automated Code Change
Google-ML-Automation
committed
Update XLA dependency to use revision http://github.com/openxla/xla/commit/c45b8fed642a7ce99e315f979dee6fc45e08f79f
Google-ML-Automation
committed
[Pallas:MGPU] Add support for indexing untiled dimensions under WG semantic.

allanrenucci
authored and
Google-ML-Automation
committed
[Mosaic GPU] Add support for indexing into splat array.

allanrenucci
authored and
Google-ML-Automation
committed
Complete reorganization of tutorials and guides
melissawm
committed
[MGPU] Add WGStridedFragLayout broadcast support on major dim.

golechwierowicz
authored and
Google-ML-Automation
committed
internal changes

al-bus
authored and
Google-ML-Automation
committed
[Mosaic GPU] Add a sanity check ensuring that layout inference only ever assigns layouts compatible with variables.

bchetioui
authored and
Google-ML-Automation
committed
Automated Code Change

cushon
authored and
Google-ML-Automation
committed
[dep] remove several dozen finalized deprecations for v0.9.0
jakevdp
committed
Remove an erroneously added license.

belitskiy
authored and
Google-ML-Automation
committed
Default make_mesh to Explicit axes for the upcoming 0.9.0 release

yashk2810
authored and
Google-ML-Automation
committed
Merge pull request #33828 from rao-ashish:asrao/cross_host_transfer_batching
Google-ML-Automation
committed
Merge pull request #33884 from melissawm:advanced-ad
Google-ML-Automation
committed
Enable using custom hermetic NCCL version.

ybaturina
authored and
Google-ML-Automation
committed
Merge pull request #34058 from jakevdp:finalize-deps
Google-ML-Automation
committed
[dep] remove deprecated `jax_safer_randint` configuration for JAX v0.9.0

Jake VanderPlas
authored and
Google-ML-Automation
committed
[dep] finalize deprecation of lax.dot positional arguments for JAX v0.9.0.

Jake VanderPlas
authored and
Google-ML-Automation
committed
[dep] remove references to the already-deprecated interpolation argument
jakevdp
committed
Remove pre-0.8.2 version guards.

belitskiy
authored and
Google-ML-Automation
committed
[dep] Remove deprecated jax.lib submodules for JAX v0.9.0.

Jake VanderPlas
authored and
Google-ML-Automation
committed
Merge pull request #34060 from jakevdp:interpolation-dep
Google-ML-Automation
committed
[hijax] upgrade VJPHiPrimitive.def_fwd to support symbolic zeros
mattjj
committed
Merge pull request #34065 from mattjj:custom-vjp3-revival
Google-ML-Automation
committed
Integrate LLVM at llvm/llvm-project@7d381f2a5634
Google-ML-Automation
committed
[dep] make jax_default_dtype_bits config a no-op for JAX v0.9.0.

Jake VanderPlas
authored and
Google-ML-Automation
committed

Commits on Dec 21, 2025

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

Commits on Dec 22, 2025

Commits on Dec 23, 2025

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

Commits on Dec 24, 2025

Add iree_metal to platforms with buffer donation support

robtaylor
and
claude
committed