Merge pull request #285 from ROCm/ci-upstream-sync-147_1

CI: 03/14/25 upstream sync

Author: github-actions[bot]

.github/workflows/pytest_cpu.yml

@@ -116,21 +116,8 @@ jobs:
           exit 1
       - name: Install Python dependencies
         run: |
-          # Remove installation of NVIDIA wheels for CPU tests.
-          sed -i 's/-r gpu-test-requirements.txt/# -r gpu-test-requirements.txt/g' build/requirements.in
-
-          # TODO(srnitin): Remove after uv is installed in the Windows Dockerfile
           $JAXCI_PYTHON -m pip install uv~=0.5.30
-          # python 3.13t cannot compile zstandard 0.23.0 due to
-          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
-          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
-          # in requirements.in.
-          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
-            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
-            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
-          else
-            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
-          fi
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/pytest_cuda.yml

@@ -100,17 +100,7 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: |
-          # python 3.13t cannot compile zstandard 0.23.0 due to
-          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
-          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
-          # in requirements.in.
-          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
-            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
-            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
-          else
-            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
-          fi
+        run: $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/wheel_tests_continuous.yml

@@ -29,6 +29,7 @@ concurrency:
 jobs:
   build-jax-artifact:
     uses: ./.github/workflows/build_artifacts.yml
+    name: "Build jax artifact"
     with:
         # Note that since jax is a pure python package, the runner OS and Python values do not
         # matter. In addition, cloning main XLA also has no effect.
@@ -46,6 +47,10 @@ jobs:
             runner: ["linux-x86-n2-16", "linux-arm64-t2a-48", "windows-x86-n2-64"]
             artifact: ["jaxlib"]
             python: ["3.10"]
+    # Note: For reasons unknown, Github actions groups jobs with the same top-level name in the
+    # dashboard only if we use an expression in the "name" field. Otherwise, it appends the matrix
+    # values to the name and creates a separate entry for each matrix combination.
+    name: "Build ${{ format('{0}', 'jaxlib') }} artifacts"
     with:
         runner: ${{ matrix.runner }}
         artifact: ${{ matrix.artifact }}
@@ -63,6 +68,7 @@ jobs:
           runner: ["linux-x86-n2-16"]
           artifact: ["jax-cuda-plugin", "jax-cuda-pjrt"]
           python: ["3.10",]
+    name: "Build ${{ format('{0}', 'CUDA') }} artifacts"
     with:
       runner: ${{ matrix.runner }}
       artifact: ${{ matrix.artifact }}
@@ -86,6 +92,7 @@ jobs:
           runner: ["linux-x86-n2-64", "linux-arm64-t2a-48", "windows-x86-n2-64"]
           python: ["3.10",]
           enable-x64: [1, 0]
+    name: "Pytest CPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -115,6 +122,7 @@ jobs:
             - runner: "linux-x86-a3-8g-h100-8gpu"
               python: "3.10"
               enable-x64: 0
+    name: "Pytest CUDA (JAX artifacts version =  ${{ format('{0}', 'head') }})"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -137,6 +145,7 @@ jobs:
           runner: ["linux-x86-g2-48-l4-4gpu",]
           python: ["3.10",]
           enable-x64: [1, 0]
+    name: "Bazel CUDA Non-RBE (JAX artifacts version =  ${{ format('{0}', 'head') }})"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -160,7 +169,7 @@ jobs:
             {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"}
           ]
-    name: "TPU tests (jax=head, jaxlib=head)"
+    name: "Pytest TPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"
     with:
       runner: ${{ matrix.tpu-specs.runner }}
       cores: ${{ matrix.tpu-specs.cores }}

.github/workflows/wheel_tests_nightly_release.yml

@@ -36,6 +36,7 @@ jobs:
           exclude:
             - runner: "windows-x86-n2-64"
               python: "3.13-nogil"
+    name: "Pytest CPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -53,6 +54,7 @@ jobs:
           python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil"]
           cuda: ["12.3", "12.1"]
           enable-x64: [0]
+    name: "Pytest CUDA (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -88,7 +90,7 @@ jobs:
             - tpu-specs:
                 type: "v5e-8"
               python: "3.11"
-    name: "TPU tests (jax=${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }}, jaxlib=${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
+    name: "Pytest TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
       runner: ${{ matrix.tpu-specs.runner }}
       cores: ${{ matrix.tpu-specs.cores }}

BUILD.bazel

@@ -31,9 +31,6 @@ transitive_py_deps(
         "//jax:experimental",
         "//jax:experimental_colocated_python",
         "//jax:experimental_sparse",
-        "//jax:internal_export_back_compat_test_util",
-        "//jax:internal_test_harnesses",
-        "//jax:internal_test_util",
         "//jax:lax_reference",
         "//jax:pallas_experimental_gpu_ops",
         "//jax:pallas_gpu_ops",

CHANGELOG.md

@@ -42,6 +42,7 @@ Patch release of 0.5.1
     {func}`jax.lax.reduce_and`, {func}`jax.lax.reduce_or`, and {func}`jax.lax.reduce_xor`.
   * {func}`jax.lax.linalg.qr`, and {func}`jax.scipy.linalg.qr`, now support
     column-pivoting on CPU and GPU. See {jax-issue}`#20282` and
+  * Added {func}`jax.random.multinomial`.
     {jax-issue}`#25955` for more details.
 
 * Changes

README.md

@@ -456,4 +456,3 @@ For details about the JAX API, see the
 
 For getting started as a JAX developer, see the
 [developer documentation](https://jax.readthedocs.io/en/latest/developer.html).
-

benchmarks/api_benchmark.py

@@ -21,14 +21,14 @@
 import google_benchmark
 import jax
 from jax import lax
-from jax._src.ad_checkpoint import checkpoint  # new jax.remat implementation
-from jax._src import core
-from jax._src.lib import xla_client as xc
 from jax._src import array
+from jax._src import core
 from jax._src import op_shardings
+from jax._src.ad_checkpoint import checkpoint  # new jax.remat implementation
+from jax._src.lib import xla_client as xc
 from jax._src.pjit import pjit_check_aval_sharding
-from jax.experimental import pjit as pjit_lib
 from jax.experimental import multihost_utils
+from jax.experimental import pjit as pjit_lib
 import jax.numpy as jnp
 import numpy as np
 
@@ -860,29 +860,44 @@ def safe_zip(state):
 
 @google_benchmark.register
 def bench_make_array_from_callback_fully_replicated_sharding(state):
-  mesh = jax.sharding.Mesh(
-      np.array(jax.devices()[:8]).reshape((4, 2)), ('x', 'y'))
-  shape = (8, 2)
-  np_arr = np.arange(16).reshape(shape)
-  s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+  mesh = create_mesh((4, 2), ('x', 'y'), state)
+  if mesh is None:
+    return
+  input_shape = (8, 2)
+  np_arr = np.arange(math.prod(input_shape)).reshape(input_shape)
 
+  s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
   while state:
-    jax.make_array_from_callback(shape, s, np_arr.__getitem__)
+    jax.make_array_from_callback(input_shape, s, np_arr.__getitem__)
 
 
 @google_benchmark.register
 @google_benchmark.option.unit(google_benchmark.kMillisecond)
-def bench_make_array_from_callback_sharded(state):
-  global_mesh = create_mesh((4, 2), ('x', 'y'), state)
+def bench_make_array_from_callback_partially_replicated_sharding(state):
+  mesh = create_mesh((4, 2), ('x', 'y'), state)
+  if mesh is None:
+    return
   input_shape = (8, 2)
-  input_data = np.arange(math.prod(input_shape)).reshape(input_shape)
+  np_arr = np.arange(math.prod(input_shape)).reshape(input_shape)
+
+  s = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None, 'y'))
+  while state:
+    jax.make_array_from_callback(input_shape, s, np_arr.__getitem__)
 
-  def callback(index):
-    return input_data[index]
 
-  s = jax.NamedSharding(global_mesh, jax.sharding.PartitionSpec('x', 'y'))
+@google_benchmark.register
+@google_benchmark.option.unit(google_benchmark.kMillisecond)
+def bench_make_array_from_callback_fully_sharded_sharding(state):
+  mesh = create_mesh((4, 2), ('x', 'y'), state)
+  if mesh is None:
+    return
+  input_shape = (8, 2)
+  np_arr = np.arange(math.prod(input_shape)).reshape(input_shape)
+
+  s = jax.NamedSharding(mesh, jax.sharding.PartitionSpec('x', 'y'))
   while state:
-    jax.make_array_from_callback((8, 2), s, callback)
+    jax.make_array_from_callback(input_shape, s, np_arr.__getitem__)
+
 
 @google_benchmark.register
 @google_benchmark.option.unit(google_benchmark.kMillisecond)

build/test-requirements.txt

@@ -7,7 +7,7 @@ flatbuffers
 hypothesis
 mpmath>=1.3
 pillow>=10.4.0
-# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t 
+# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
 portpicker; python_version<"3.13"
 pytest-xdist
 wheel
@@ -19,3 +19,6 @@ matplotlib~=3.8.4; python_version=="3.10"
 matplotlib; python_version>="3.11"
 opt-einsum
 auditwheel
+
+# CPU Pytests crash with NumPy 2.2+ on Linux Aarch64; b/399168632
+numpy~=2.1.0; platform_system == "Linux" and platform_machine == "aarch64"

docs/_static/style.css

@@ -296,3 +296,13 @@ html[data-theme="light"] .highlight span.gt {
 html[data-theme="light"] .highlight span.gr {
   color: #ff0000;
 }
+
+.header-article-items__start {
+  display: flex;
+  flex-direction: row;
+  gap: 0.5em;
+}
+
+.bd-breadcrumbs {
+  margin-bottom: 0;
+}