Add iree_metal to platforms with buffer donation support

.bazelrc

@@ -178,20 +178,21 @@ common:clang --copt=-Wno-error=c23-extensions
 common:cuda_v12 --repo_env=HERMETIC_CUDA_VERSION="12.9.1"
 common:cuda_v12 --repo_env=HERMETIC_CUDNN_VERSION="9.8.0"
 common:cuda_v12 --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.9"
+common:cuda_v12 --repo_env=HERMETIC_NCCL_VERSION="2.27.7"
 # "sm" means we emit only cubin, which is forward compatible within a GPU generation.
 # "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
 common:cuda_v12 --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,sm_90,sm_100,compute_120"
 
 common:cuda_v13 --repo_env=HERMETIC_CUDA_VERSION="13.0.0"
 common:cuda_v13 --repo_env=HERMETIC_CUDNN_VERSION="9.12.0"
 common:cuda_v13 --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.20"
+common:cuda_v13 --repo_env=HERMETIC_NCCL_VERSION="2.27.7"
 common:cuda_v13 --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_75,sm_80,sm_90,sm_100,compute_120"
 
 common:cuda_common --repo_env TF_NEED_CUDA=1
 common:cuda_common --repo_env TF_NCCL_USE_STUB=1
 common:cuda_common --@local_config_cuda//:enable_cuda
 common:cuda_common --@local_config_cuda//cuda:include_cuda_libs=true
-common:cuda_common --@cuda_driver//:include_cuda_umd_libs=true
 
 # Force the linker to set RPATH, not RUNPATH. When resolving dynamic libraries,
 # ld.so prefers in order: RPATH, LD_LIBRARY_PATH, RUNPATH. JAX sets RPATH to
@@ -218,7 +219,8 @@ common:cuda --config=cuda12
 
 # This config is used for building targets with CUDA/NVSHMEM libraries from stubs.
 common:cuda_libraries_from_stubs --@local_config_cuda//cuda:include_cuda_libs=false
-common:cuda_libraries_from_stubs --@cuda_driver//:include_cuda_umd_libs=false
+
+common:hermetic_cuda_umd --@cuda_driver//:include_cuda_umd_libs=true
 
 # common CUDA and other C++ targets with Clang
 common:build_cuda_with_clang --@local_config_cuda//:cuda_compiler=clang
@@ -396,6 +398,9 @@ common:rbe --spawn_strategy=remote,worker,standalone,local
 common:rbe --remote_download_toplevel
 test:rbe --test_env=USER=anon
 
+common:rbe_cpu_pool --repo_env=REMOTE_GPU_TESTING=0
+common:rbe_gpu_pool --repo_env=REMOTE_GPU_TESTING=1
+
 # RBE configs for Linux x86
 # Set the remote worker pool
 common:rbe_linux_x86_64_base --remote_instance_name=projects/tensorflow-testing/instances/default_instance
@@ -416,7 +421,9 @@ common:rbe_linux_x86_64 --config=rbe_linux_x86_64_base
 common:rbe_linux_x86_64 --config=ci_linux_x86_64
 
 common:rbe_linux_x86_64_cuda_common --config=rbe_linux_x86_64_base
-common:rbe_linux_x86_64_cuda_common --repo_env=REMOTE_GPU_TESTING=1
+common:rbe_linux_x86_64_cuda_common --config=rbe_gpu_pool
+# Update UMD version when RBE CUDA driver is updated.
+common:rbe_linux_x86_64_cuda_common --repo_env=HERMETIC_CUDA_UMD_VERSION="13.0.1"
 
 common:rbe_linux_x86_64_cuda12 --config=rbe_linux_x86_64_cuda_common
 common:rbe_linux_x86_64_cuda12 --config=ci_linux_x86_64_cuda12

.github/actionlint.yaml

@@ -16,4 +16,5 @@ self-hosted-runner:
     - "linux-x86-ct6e-180-8tpu" # Linux X86 TPU runner using ct6e-hightpu-8t machine with 2x4 topology.
     - "linux-x86-ct6e-180-4tpu" # Linux X86 TPU runner using ct6e-hightpu-4t machine with 2x2 topology.
     - "linux-x86-ct4p-240-4tpu" # Linux X86 TPU runner using ct4p-hightpu-4t machine with 2x2x1 topology.
+    - "linux-x86-tpu7x-224-4tpu" # Linux X86 TPU runner using tpu7x-224 machine with 4 TPU chips (8 cores) and 2x2x1 topology.
     - "linux-x86_64-cirrascale-64-8gpu-amd-mi250" # AMD runner

.github/workflows/bazel_cuda_h100_b200.yml

@@ -41,10 +41,10 @@ jobs:
         uses: tj-actions/changed-files@ed68ef82c095e0d48ec87eccea555d944a631a4c # v46
         with:
           files: |
-            jax/jax/_src/pallas/mosaic_gpu/**
-            jax/jax/experimental/mosaic/gpu/**
-            jax/jaxlib/mosaic/dialect/gpu/**
-            jax/jaxlib/mosaic/gpu/**
+            jax/_src/pallas/mosaic_gpu/**
+            jax/experimental/mosaic/gpu/**
+            jaxlib/mosaic/dialect/gpu/**
+            jaxlib/mosaic/gpu/**
             .github/workflows/bazel_cuda_h100_b200.yml
       - name: List all changed files
         env:
@@ -76,6 +76,7 @@ jobs:
             bazel test \
             --config=ci_linux_x86_64_cuda \
             --config=ci_rbe_cache \
+            --config=hermetic_cuda_umd \
             --repo_env=HERMETIC_PYTHON_VERSION="3.14" \
             --repo_env=HERMETIC_CUDNN_VERSION="9.11.0" \
             --repo_env=HERMETIC_CUDA_UMD_VERSION="13.0.0" \
@@ -120,6 +121,7 @@ jobs:
             bazel test \
             --config=ci_linux_x86_64_cuda \
             --config=ci_rbe_cache \
+            --config=hermetic_cuda_umd \
             --repo_env=HERMETIC_PYTHON_VERSION="3.14" \
             --repo_env=HERMETIC_CUDNN_VERSION="9.11.0" \
             --repo_env=HERMETIC_CUDA_UMD_VERSION="13.0.0" \

.github/workflows/bazel_cuda_presubmit.yml

@@ -73,6 +73,9 @@ jobs:
             - python: "3.14"
               enable-x64: 1
               jaxlib-version: "pypi_latest"
+            # Exclude CUDA 12 on jaxlib head because it's too slow.
+            - cuda-version: "12"
+              jaxlib-version: "head"
     name: "Bazel single accelerator ${{ format('{0}', 'CUDA tests') }}"
 # End Presubmit Naming Check github-cuda-presubmits
     with:

.github/workflows/build_artifacts.yml

@@ -115,7 +115,7 @@ jobs:
     name: "${{ inputs.artifact }},
             ${{ (contains(inputs.runner, 'linux-x86') && 'linux x86') ||
             (contains(inputs.runner, 'linux-arm64') && 'linux arm64') ||
-            (contains(inputs.runner, 'windows-x86') && 'windows x86') }}, py ${{ inputs.python }} ${{ (contains(inputs.artifact, 'cuda') && format(', cuda {0}', inputs.cuda-version)) || '' }}, clone main XLA=${{ inputs.clone_main_xla }}"
+            (contains(inputs.runner, 'windows-x86') && 'windows x86') }}, py ${{ inputs.python }}${{ (contains(inputs.artifact, 'cuda') && format(', cuda {0}', inputs.cuda-version)) || '' }}, clone main XLA=${{ inputs.clone_main_xla }}"
 
     # Map the job outputs to step outputs
     outputs:

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -35,7 +35,8 @@ jobs:
         tpu: [
           {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
           {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-          {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
+          {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"},
+          {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"}
         ]
         python-version: ["3.11"]
         # Exclude v6e-8 tests for pypi_latest for resource constraints.

.github/workflows/wheel_tests_continuous.yml

@@ -241,7 +241,8 @@ jobs:
           tpu-specs: [
             # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
+            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"},
+            {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"}
           ]
           libtpu-version-type: ["nightly"]
     name: "Pytest TPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"
@@ -267,6 +268,7 @@ jobs:
           tpu-specs: [
             {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
+            {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"},
           ]
           libtpu-version-type: ["nightly"]
     name: "Bazel tests TPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"

.github/workflows/wheel_tests_nightly_release.yml

@@ -177,7 +177,8 @@ jobs:
           tpu-specs: [
             # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
+            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"},
+            {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"}
           ]
           libtpu-version-type: ["pypi_latest", "nightly"]
           exclude:
@@ -192,13 +193,17 @@ jobs:
             - tpu-specs:
                 type: "v6e-8"
               python: "3.13-nogil"
-            # Run min and max Python versions for v5e-8
+            # Run max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
               python: "3.11"
             - tpu-specs:
                 type: "v5e-8"
               python: "3.12"
+            # Run min and max Python versions for v7x-8
+            - tpu-specs:
+                type: "v7x-8"
+              python: "3.12"
 
     name: "Pytest TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
@@ -222,6 +227,7 @@ jobs:
             # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
             {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
+            {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"},
           ]
           libtpu-version-type: ["pypi_latest", "nightly"]
           exclude:
@@ -239,12 +245,13 @@ jobs:
             - tpu-specs:
                 type: "v4-8"
               python: "3.13-nogil"
-            # Run min and max Python versions for v5e-8
+            # Run max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
-              python: "3.11"
+              python: "3.12"
+            # Run min and max Python versions for v7x-8
             - tpu-specs:
-                type: "v5e-8"
+                type: "v7x-8"
               python: "3.12"
 
     name: "Bazel tests TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"

.pre-commit-config.yaml

@@ -41,7 +41,7 @@ repos:
   - id: mypy
     files: (jax/|tests/typing_test\.py)
     exclude: jax/_src/basearray.py|jax/numpy/__init__.py|jax/nn/__init__.py|jaxlib/_jax/.*  # Use pyi instead
-    additional_dependencies: [types-requests==2.31.0, numpy>=2.2.0, scipy-stubs]
+    additional_dependencies: [types-requests==2.31.0, numpy~=2.3.0, scipy-stubs]
     args: [--config=pyproject.toml]
 
 - repo: https://github.com/mwouts/jupytext

BUILD.bazel

@@ -30,20 +30,43 @@ wheel_sources(
     data_srcs = ["//jax"],
     py_srcs = [
         "//jax",
-        "//jax:compilation_cache",
-        "//jax:experimental",
-        "//jax:experimental_colocated_python",
-        "//jax:experimental_sparse",
-        "//jax:experimental_buffer_callback",
-        "//jax:experimental_serialize_executable",
-        "//jax:pallas_experimental_gpu_ops",
-        "//jax:pallas_fuser",
-        "//jax:pallas_gpu_ops",
-        "//jax:pallas_mosaic_gpu",
-        "//jax:pallas_tpu_ops",
-        "//jax:pallas_triton",
-        "//jax:source_mapper",
-        "//jax:sparse_test_util",
+        "//jax/example_libraries:example_libraries",
+        "//jax/example_libraries:optimizers",
+        "//jax/example_libraries:stax",
+        "//jax/experimental:buffer_callback",
+        "//jax/experimental:checkify",
+        "//jax/experimental:colocated_python",
+        "//jax/experimental:compilation_cache",
+        "//jax/experimental:compute_on",
+        "//jax/experimental:custom_dce",
+        "//jax/experimental:custom_partitioning",
+        "//jax/experimental:fused",
+        "//jax/experimental:hijax",
+        "//jax/experimental:jet",
+        "//jax/experimental:layout",
+        "//jax/experimental:mesh_utils",
+        "//jax/experimental:multihost_utils",
+        "//jax/experimental:ode",
+        "//jax/experimental:pallas_experimental_gpu_ops",
+        "//jax/experimental:pallas_fuser",
+        "//jax/experimental:pallas_gpu_ops",
+        "//jax/experimental:pallas_mosaic_gpu",
+        "//jax/experimental:pallas_tpu_ops",
+        "//jax/experimental:pallas_triton",
+        "//jax/experimental:pjit",
+        "//jax/experimental:profiler",
+        "//jax/experimental:rnn",
+        "//jax/experimental:scheduling_groups",
+        "//jax/experimental:serialize_executable",
+        "//jax/experimental:shard_alike",
+        "//jax/experimental:shard_map",
+        "//jax/experimental:source_mapper",
+        "//jax/experimental:sparse_test_util",
+        "//jax/experimental:sparse",
+        "//jax/experimental:topologies",
+        "//jax/experimental:transfer",
+        "//jax/experimental:xla_metadata",
+        "//jax/experimental",
         "//jax/_src:lax_reference",
         "//jax/_src:internal_export_back_compat_test_util",
         "//jax/_src:internal_export_back_compat_test_data",
@@ -123,6 +146,7 @@ genrule(
         "//jax/experimental/mosaic/gpu/examples:flash_attention.py",
         "//jax/experimental/mosaic/gpu/examples:matmul.py",
         "//jax/_src:test_multiprocess",
+        "//jax/_src/pallas:pallas_test_util",
     ],
     outs = ["wheel_additives.zip"],
     cmd = "$(location @bazel_tools//tools/zip:zipper) c $@ $(SRCS)",

CHANGELOG.md

@@ -16,9 +16,19 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 
 ## Unreleased
 
+## JAX 0.8.2 (December 18, 2025)
+
 * Deprecations
   * `jax.lax.pvary` has been deprecated.
     Please use `jax.lax.pcast(..., to='varying')` as the replacement.
+  * Complex arguments passed to {func}`jax.numpy.arange` now result in a
+    deprecation warning, because the output is poorly-defined.
+  * From {mod}`jax.core` a number of symbols are newly deprecated including:
+    `call_impl`, `get_aval`, `mapped_aval`, `subjaxprs`, `set_current_trace`,
+    `take_current_trace`, `traverse_jaxpr_params`, `unmapped_aval`,
+    `AbstractToken`,  and `TraceTag`.
+  * All symbols in {mod}`jax.interpreters.pxla` are deprecated. These are
+    primarily JAX internal APIs, and users should not rely on them.
 
 * Changes:
   * jax's `Tracer` no longer inherits from `jax.Array` at runtime. However,
@@ -29,6 +39,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     For the moment, during Python type checking, we continue to declare `Tracer`
     as a subclass of `Array`, however we expect to remove this in a future
     release.
+  * `jax.experimental.si_vjp` has been deleted.
+    `jax.vjp` subsumes it's functionality.
 
 ## JAX 0.8.1 (November 18, 2025)
 

WORKSPACE

@@ -1,6 +1,7 @@
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
 # The XLA commit is determined by third_party/xla/revision.bzl.
 load("//third_party/xla:workspace.bzl", jax_xla_workspace = "repo")
-load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 jax_xla_workspace()
 
@@ -12,15 +13,15 @@ load("@xla//:workspace3.bzl", "xla_workspace3")
 
 xla_workspace3()
 
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-
 # Initialize Hermetic toolchains
 # Details: https://github.com/google-ml-infra/rules_ml_toolchain
 tf_http_archive(
     name = "rules_ml_toolchain",
-    sha256 = "8123d826b0a4c5ceda767abc8092419fcc980c3ce45fb0f438b101fb886c014c",
-    strip_prefix = "rules_ml_toolchain-552b53a04a86fd5cdb4d5091e7420411d8b2a045",
-    urls = tf_mirror_urls("https://github.com/google-ml-infra/rules_ml_toolchain/archive/552b53a04a86fd5cdb4d5091e7420411d8b2a045.tar.gz"),
+    sha256 = "1c2c530a054e9e8b3c811ec21ed8a687fc865bec3abbc8ff65beb829b1d67ae4",
+    strip_prefix = "rules_ml_toolchain-6734d2a174bf29e731d3f473743d1cc1a86100c3",
+    urls = tf_mirror_urls(
+        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/6734d2a174bf29e731d3f473743d1cc1a86100c3.tar.gz",
+    ),
 )
 
 load(

benchmarks/mosaic/BUILD

@@ -35,7 +35,7 @@ jax_multiplatform_test(
     enable_configs = ["gpu_h100"],
     tags = ["notap"],
     deps = [
-        "//jax:mosaic_gpu",
+        "//jax/experimental:mosaic_gpu",
         "//jax/experimental/mosaic/gpu/examples:matmul",
         "//third_party/py/google_benchmark",
     ] + py_deps("absl/testing") + py_deps("numpy"),

build/requirements.in

@@ -18,7 +18,7 @@ wheel
 # the requirements files.
 jaxlib
 
-jax-cuda12-plugin; sys_platform == "linux" and python_version<"3.14"
+jax-cuda12-plugin; sys_platform == "linux"
 jax-cuda13-plugin
 jax-cuda12-pjrt; sys_platform == "linux"
 jax-cuda13-pjrt