Merge branch 'main' of https://github.com/jax-ml/jax

Author: nitins17

.github/workflows/asan.yaml

@@ -12,7 +12,7 @@ on:
     branches:
       - main
     paths:
-      - '**/workflows/asan.yml'
+      - '**/workflows/asan.yaml'
 
 jobs:
   asan:

.github/workflows/bazel_cpu_rbe.yml

@@ -1,4 +1,4 @@
-name: Run Bazel CPU tests (RBE)
+name: CI - Bazel CPU tests (RBE)
 
 on:
   # pull_request:
@@ -15,32 +15,31 @@ on:
         - 'yes'
         - 'no'
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
 jobs:
-  run_bazel_rbe_cpu_tests:
-    continue-on-error: true
-    defaults:
-      run:
-        # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
-        shell: bash
+  run_tests:
+    if: github.event.repository.fork == false
     strategy:
       matrix:
         runner: ["windows-x86-n2-64", "linux-x86-n2-16", "linux-arm64-t2a-16"]
 
     runs-on: ${{ matrix.runner }}
-    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build:670606426-python3.12') ||
-                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
+    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
+                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') ||
                    (contains(matrix.runner, 'windows-x86') && null) }}
 
+
     env:
-      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
 
     steps:
       - uses: actions/checkout@v3
-      # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Run Bazel CPU Tests
+      - name: Run Bazel CPU Tests with RBE
         run: ./ci/run_bazel_test_cpu_rbe.sh

.github/workflows/bazel_gpu_rbe.yml

@@ -1,4 +1,4 @@
-name: Run Bazel GPU tests (RBE)
+name: CI - Bazel GPU tests (RBE)
 
 on:
   # pull_request:
@@ -15,26 +15,28 @@ on:
         - 'yes'
         - 'no'
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
 jobs:
-  build:
+  run_tests:
+    if: github.event.repository.fork == false
     strategy:
       matrix:
         runner: ["linux-x86-n2-16"]
 
     runs-on: ${{ matrix.runner }}
-    container:
-      image: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+    container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
 
     env:
-      JAXCI_CLONE_MAIN_XLA: 1
-      JAXCI_HERMETIC_PYTHON_VERSION: 3.12
+      JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
 
     steps:
       - uses: actions/checkout@v3
-      # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Run Bazel GPU tests using RBE
+      - name: Run Bazel GPU Tests with RBE
         run: ./ci/run_bazel_test_gpu_rbe.sh

.github/workflows/ci-build.yaml

@@ -35,7 +35,7 @@ jobs:
         with:
           python-version: 3.11
       - run: python -m pip install pre-commit
-      - uses: actions/cache@3624ceb22c1c5a301c8db4169662070a689d9ea8 # v4.1.1
+      - uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2
         with:
           path: ~/.cache/pre-commit
           key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('.pre-commit-config.yaml', 'setup.py') }}

.github/workflows/jax-array-api.yml

@@ -28,7 +28,7 @@ jobs:
       with:
         repository: data-apis/array-api-tests
         # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'bcd5919bbbdf4d4806b5b2613b4d8c0bc0625c54'  # Latest commit as of 2024-10-31 👻
+        ref: 'a3f3f376308e64f0ac15b307dfe27be945409e41'  # Latest commit as of 2024-11-14
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python ${{ matrix.python-version }}

CHANGELOG.md

@@ -40,6 +40,13 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
         `platforms` instead.
   * Hashing of tracers, which has been deprecated since version 0.4.30, now
     results in a `TypeError`.
+  * {func}`jax.scipy.linalg.toeplitz` now does implicit batching on multi-dimensional
+    inputs. To recover the previous behavior, you can call {func}`jax.numpy.ravel`
+    on the function inputs.
+  * {func}`jax.scipy.special.gamma` and {func}`jax.scipy.special.gammasgn` now
+    return NaN for negative integer inputs, to match the behavior of SciPy from
+    https://github.com/scipy/scipy/pull/21827.
+  * `jax.clear_backends` was removed after being deprecated in v0.4.26.
 
 * New Features
   * {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
@@ -48,6 +55,21 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.tree_util.register_dataclass` now allows metadata fields to be
     declared inline via {func}`dataclasses.field`. See the function documentation
     for examples.
+  * Added {func}`jax.numpy.put_along_axis`.
+  * {func}`jax.lax.linalg.eig` and the related `jax.numpy` functions
+    ({func}`jax.numpy.linalg.eig` and {func}`jax.numpy.linalg.eigvals`) are now
+    supported on GPU. See {jax-issue}`#24663` for more details.
+
+* Bug fixes
+  * Fixed a bug where the GPU implementations of LU and QR decomposition would
+    result in an indexing overflow for batch sizes close to int32 max. See
+    {jax-issue}`#24843` for more details.
+
+* Deprecations
+  * `jax.lib.xla_extension.ArrayImpl` and `jax.lib.xla_client.ArrayImpl` are deprecated;
+    use `jax.Array` instead.
+  * `jax.lib.xla_extension.XlaRuntimeError` is deprecated; use `jax.errors.JaxRuntimeError`
+    instead.
 
 ## jax 0.4.35 (Oct 22, 2024)
 
@@ -79,7 +101,7 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * The semi-public API `jax.lib.xla_client.register_custom_call_target` has
     been deprecated. Use the JAX FFI instead.
   * The semi-public APIs `jax.lib.xla_client.dtype_to_etype`,
-    `jax.lib.xla_client.ops`, 
+    `jax.lib.xla_client.ops`,
     `jax.lib.xla_client.shape_from_pyval`, `jax.lib.xla_client.PrimitiveType`,
     `jax.lib.xla_client.Shape`, `jax.lib.xla_client.XlaBuilder`, and
     `jax.lib.xla_client.XlaComputation` have been deprecated. Use StableHLO

README.md

@@ -193,8 +193,7 @@ You can mix `jit` and `grad` and any other JAX transformation however you like.
 
 Using `jit` puts constraints on the kind of Python control flow
 the function can use; see
-the [Gotchas
-Notebook](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#python-control-flow-+-JIT)
+the tutorial on [Control Flow and Logical Operators with JIT](https://jax.readthedocs.io/en/latest/control-flow.html)
 for more.
 
 ### Auto-vectorization with `vmap`
@@ -353,7 +352,7 @@ Some standouts:
 1. [In-place mutating updates of
    arrays](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates), like `x[i] += y`, aren't supported, but [there are functional alternatives](https://jax.readthedocs.io/en/latest/jax.ops.html). Under a `jit`, those functional alternatives will reuse buffers in-place automatically.
 1. [Random numbers are
-   different](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers), but for [good reasons](https://github.com/jax-ml/jax/blob/main/docs/jep/263-prng.md).
+   different](https://jax.readthedocs.io/en/latest/random-numbers.html), but for [good reasons](https://github.com/jax-ml/jax/blob/main/docs/jep/263-prng.md).
 1. If you're looking for [convolution
    operators](https://jax.readthedocs.io/en/latest/notebooks/convolutions.html),
    they're in the `jax.lax` package.
@@ -373,7 +372,7 @@ Some standouts:
    and NumPy types aren't preserved, namely `np.add(1, np.array([2],
    np.float32)).dtype` is `float64` rather than `float32`.
 1. Some transformations, like `jit`, [constrain how you can use Python control
-   flow](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#control-flow).
+   flow](https://jax.readthedocs.io/en/latest/control-flow.html).
    You'll always get loud errors if something goes wrong. You might have to use
    [`jit`'s `static_argnums`
    parameter](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit),
@@ -394,6 +393,7 @@ Some standouts:
 | Google TPU | yes          | n/a           | n/a          | n/a          | n/a            | n/a                 |
 | AMD GPU    | yes          | no            | experimental | n/a          | no             | no                  |
 | Apple GPU  | n/a          | no            | n/a          | experimental | n/a            | n/a                 |
+| Intel GPU  | experimental | n/a           | n/a          | n/a          | no             | no                  |
 
 
 ### Instructions
@@ -405,6 +405,7 @@ Some standouts:
 | Google TPU      | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`                 |
 | AMD GPU (Linux) | Use [Docker](https://hub.docker.com/r/rocm/jax-community/tags), [pre-built wheels](https://github.com/ROCm/jax/releases), or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
 | Mac GPU         | Follow [Apple's instructions](https://developer.apple.com/metal/jax/).                                          |
+| Intel GPU       | Follow [Intel's instructions](https://github.com/intel/intel-extension-for-openxla/blob/main/docs/acc_jax.md).  |
 
 See [the documentation](https://jax.readthedocs.io/en/latest/installation.html)
 for information on alternative installation strategies. These include compiling

ci/README.md

@@ -1,6 +1,7 @@
-# JAX continuous integration..
+# JAX continuous integration
 
-> **Warning** This folder is still under construction. It is part of an ongoing
+> [!WARNING]
+> This folder is still under construction. It is part of an ongoing
 > effort to improve the structure of CI and build related files within the
 > JAX repo. This warning will be removed when the contents of this
 > directory are stable and appropriate documentation around its usage is in

ci/envs/default.env

@@ -13,7 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 # This file contains all the default values for the "JAXCI_" environment
-# variables used in the CI scripts.
+# variables used in the CI scripts. These variables are used to control the
+# behavior of the CI scripts such as the Python version used, path to JAX/XLA
+# repo, if to clone XLA repo, etc.
 
 # The path to the JAX git repository.
 export JAXCI_JAX_GIT_DIR=$(pwd)
@@ -22,8 +24,9 @@ export JAXCI_JAX_GIT_DIR=$(pwd)
 # set.
 export JAXCI_HERMETIC_PYTHON_VERSION=${JAXCI_HERMETIC_PYTHON_VERSION:-$(python3 -V | awk '{print $2}' | awk -F. '{print $1"."$2}')}
 
-# Set JAXCI_XLA_GIT_DIR to the root of the XLA git repository if you want to
-# use a local copy of XLA instead of the pinned version in the WORKSPACE.
+# Set JAXCI_XLA_GIT_DIR to the root of the XLA git repository to use a local
+# copy of XLA instead of the pinned version in the WORKSPACE. When
+# JAXCI_CLONE_MAIN_XLA=1, this gets set automatically.
 export JAXCI_XLA_GIT_DIR=${JAXCI_XLA_GIT_DIR:-}
 
 # If set to 1, the builds will clone the XLA repository at HEAD and set its

ci/run_bazel_test_cpu_rbe.sh

@@ -22,7 +22,7 @@
 # -o allexport: export all functions and variables to be available to subscripts
 set -exu -o history -o allexport
 
-# Inherit default JAXCI environment variables.
+# Source default JAXCI environment variables.
 source ci/envs/default.env
 
 # Clone XLA at HEAD if path to local XLA is not provided
@@ -33,6 +33,7 @@ fi
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
+# Run Bazel CPU tests with RBE.
 os=$(uname -s | awk '{print tolower($0)}')
 arch=$(uname -m)
 
@@ -43,8 +44,8 @@ arch=$(uname -m)
 # single machine can take a long time, we skip running them on these
 # platforms.
 if [[ $os == "darwin" ]] || ( [[ $os == "linux" ]] && [[ $arch == "aarch64" ]] ); then
-    echo "Building RBE CPU tests..."
-    bazel build --config=rbe_cross_compile_${os}_${arch} \
+      echo "Building RBE CPU tests..."
+      bazel build --config=rbe_cross_compile_${os}_${arch} \
             --repo_env=HERMETIC_PYTHON_VERSION="$JAXCI_HERMETIC_PYTHON_VERSION" \
             --override_repository=xla="${JAXCI_XLA_GIT_DIR}" \
             --test_env=JAX_NUM_GENERATED_CASES=25 \
@@ -54,8 +55,8 @@ if [[ $os == "darwin" ]] || ( [[ $os == "linux" ]] && [[ $arch == "aarch64" ]] )
             --color=yes \
             //tests:cpu_tests //tests:backend_independent_tests
 else
-    echo "Running RBE CPU tests..."
-    bazel test --config=rbe_${os}_${arch} \
+      echo "Running RBE CPU tests..."
+      bazel test --config=rbe_${os}_${arch} \
             --repo_env=HERMETIC_PYTHON_VERSION="$JAXCI_HERMETIC_PYTHON_VERSION" \
             --override_repository=xla="${JAXCI_XLA_GIT_DIR}" \
             --test_env=JAX_NUM_GENERATED_CASES=25 \