Repository: google/jax
Pull request #20768 (2024-04-15 23:32:32)

Add more imports to jax.extend.core


Patch:
diff --git a/.github/workflows/jax-array-api.yml b/.github/workflows/jax-array-api.yml
index 01942e524..ef80582e1 100644
--- a/.github/workflows/jax-array-api.yml
+++ b/.github/workflows/jax-array-api.yml
@@ -25,7 +25,7 @@ jobs:
       with:
         repository: data-apis/array-api-tests
         # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'e38ce3466e596ed2b8fa4638f161f5563ded81a8'  # Latest commit as of 2024-04-15
+        ref: '58113931e9959ee70dca4bace430b55b90a1eed8'  # Latest commit as of 2024-03-18
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python ${{ matrix.python-version }}
@@ -42,4 +42,4 @@ jobs:
         JAX_ENABLE_X64: 'true'
       run: |
         cd ${GITHUB_WORKSPACE}/array-api-tests
-        pytest array_api_tests --max-examples=5 --derandomize --disable-deadline --skips-file ${GITHUB_WORKSPACE}/jax/experimental/array_api/skips.txt
+        pytest --ci array_api_tests --max-examples=5 --derandomize --disable-deadline --skips-file ${GITHUB_WORKSPACE}/jax/experimental/array_api/skips.txt
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a30f333d9..331554c67 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,9 +9,8 @@ Remember to align the itemized text with the first line of an item within a list
 ## jax 0.4.27
 
 * New Functionality
-  * Added {func}`jax.numpy.unstack` and {func}`jax.numpy.cumulative_sum`,
-    following their addition in the array API 2023 standard, soon to be
-    adopted by NumPy.
+  * Added {func}`jax.numpy.unstack`, following the addition of this function in
+    the array API 2023 standard, soon to be adopted by NumPy.
 
 * Changes
   * {func}`jax.pure_callback` and {func}`jax.experimental.io_callback`
diff --git a/docs/conf.py b/docs/conf.py
index 81784cd93..5b7bf00d7 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -126,7 +126,6 @@ exclude_patterns = [
     'jep/9407-type-promotion.md',
     'jax-101/*.md',
     'autodidax.md',
-    'tutorials/sharded-computation.md',
 ]
 
 # The name of the Pygments (syntax highlighting) style to use.
@@ -214,7 +213,6 @@ nb_execution_excludepatterns = [
     # Requires accelerators
     'pallas/quickstart.*',
     'pallas/tpu/pipelining.*',
-    'tutorials/sharded-computation.*'
 ]
 
 # -- Options for HTMLHelp output ---------------------------------------------
diff --git a/docs/installation.md b/docs/installation.md
index 22b0d1aa9..146694686 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -1,50 +1,43 @@
-(installation)=
 # Installing JAX
 
-Using JAX requires installing two packages: `jax`, which is pure Python and
-cross-platform, and `jaxlib` which contains compiled binaries, and requires
-different builds for different operating systems and accelerators.
+JAX is written in pure Python, but it depends on XLA, which needs to be
+installed as the `jaxlib` package. Use the following instructions to install a
+binary package with `pip` or `conda`, to use a
+[Docker container](#docker-containers-nvidia-gpu), or to [build JAX from
+source](developer.md#building-from-source).
 
-**TL;DR** For most users, a typical JAX installation may look something like this:
-
-* **CPU-only (Linux/macOS/Windows)**
-  ```
-  pip install -U "jax[cpu]"
-  ```
-* **GPU (NVIDIA, CUDA 12, x86_64)**
-  ```
-  pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-  ```
-
-(install-supported-platforms)=
 ## Supported platforms
 
-The table below shows all supported platforms and installation options. Check if your setup is supported; and if it says _"yes"_ or _"experimental"_, then click on the corresponding link to learn how to install JAX in greater detail.
+|            | Linux x86_64 | Linux aarch64           | Mac x86_64   | Mac ARM        | Windows x86_64 | Windows WSL2 x86_64 |
+|------------|--------------|-------------------------|--------------|----------------|----------------|---------------------|
+| CPU        | [yes](#cpu)         | [yes](#cpu) | [yes](#cpu)          | [yes](#cpu)            | [yes](#cpu)            | [yes](#cpu)                 |
+| NVIDIA GPU | [yes](#nvidia-gpu)               | [yes](#nvidia-gpu) | no           | n/a            | no             | [experimental](#nvidia-gpu)        |
+| Google TPU | [yes](#google-tpu)  | n/a                     | n/a          | n/a            | n/a            | n/a                 |
+| AMD GPU    | [experimental](#amd-gpu) | no                      | no           | n/a                 | no             | no                  |
+| Apple GPU  | n/a                 | no                      | [experimental](#apple-gpu) | [experimental](#apple-gpu)   | n/a            | n/a                 |
 
-|                  | Linux, x86_64                        | Linux, aarch64      | macOS, Intel x86_64, AMD GPU   | macOS, Apple Silicon, ARM-based       | Windows, x86_64         | Windows WSL2, x86_64           |
-|------------------|---------------------------------------|--------------------------------|----------------------------------------|----------------------------------------|-------------------------|-----------------------------------------|
-| CPU              | {ref}`yes <install-cpu>`              | {ref}`yes <install-cpu>`        | {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>`|
-| NVIDIA GPU       | {ref}`yes <install-nvidia-gpu>`       | {ref}`yes <install-nvidia-gpu>` | no | n/a | no | {ref}`experimental <install-nvidia-gpu>` |
-| Google Cloud TPU | {ref}`yes <install-google-tpu>`       | n/a | n/a | n/a | n/a | n/a |
-| AMD GPU          | {ref}`experimental <install-amd-gpu>` | no | no | n/a | no | no |
-| Apple GPU    | n/a                                   | no | {ref}`experimental <install-apple-gpu>` | {ref}`experimental <install-apple-gpu>` | n/a |  n/a |
 
+We support installing or building `jaxlib` on Linux (Ubuntu 20.04 or later) and
+macOS (10.12 or later) platforms. There is also *experimental* native Windows
+support.
+
+Windows users can use JAX on CPU and GPU via the [Windows Subsystem for
+Linux](https://docs.microsoft.com/en-us/windows/wsl/about), or alternatively
+they can use the native Windows CPU-only support.
 
-(install-cpu)=
 ## CPU
 
 ### pip installation: CPU
 
-Currently, the JAX team releases `jaxlib` wheels for the following
+We currently release `jaxlib` wheels for the following
 operating systems and architectures:
-
-- Linux, x86_64
-- macOS, Intel
-- macOS, Apple ARM-based
-- Windows, x86_64 (*experimental*)
+* Linux, x86-64
+* Mac, Intel
+* Mac, ARM
+* Windows, x86-64 (*experimental*)
 
 To install a CPU-only version of JAX, which might be useful for doing local
-development on a laptop, you can run:
+development on a laptop, you can run
 
 ```bash
 pip install --upgrade pip
@@ -60,65 +53,56 @@ to pip install on other operating systems and architectures may lead to `jaxlib`
 not being installed alongside `jax`, although `jax` may successfully install
 (but fail at runtime).
 
-
-(install-nvidia-gpu)=
 ## NVIDIA GPU
 
 JAX supports NVIDIA GPUs that have SM version 5.2 (Maxwell) or newer.
 Note that Kepler-series GPUs are no longer supported by JAX since
 NVIDIA has dropped support for Kepler GPUs in its software.
 
-You must first install the NVIDIA driver. You're
-recommended to install the newest driver available from NVIDIA, but the driver
-version must be >= 525.60.13 for CUDA 12 on Linux.
-
+You must first install the NVIDIA driver. We
+recommend installing the newest driver available from NVIDIA, but the driver
+must be version >= 525.60.13 for CUDA 12 on Linux.
 If you need to use a newer CUDA toolkit with an older driver, for example
 on a cluster where you cannot update the NVIDIA driver easily, you may be
 able to use the
 [CUDA forward compatibility packages](https://docs.nvidia.com/deploy/cuda-compatibility/)
 that NVIDIA provides for this purpose.
 
-### pip installation: NVIDIA GPU (CUDA, installed via pip, easier)
-
-There are two ways to install JAX with NVIDIA GPU support:
+### pip installation: GPU (CUDA, installed via pip, easier)
 
-- Using NVIDIA CUDA and cuDNN installed from pip wheels
-- Using a self-installed CUDA/cuDNN
-
-The JAX team strongly recommends installing CUDA and cuDNN using the pip wheels,
-since it is much easier!
-
-This method is only supported on x86_64, because NVIDIA has not released aarch64
-CUDA pip packages.
+There are two ways to install JAX with NVIDIA GPU support: using CUDA and CUDNN
+installed from pip wheels, and using a self-installed CUDA/CUDNN. We
+strongly recommend installing CUDA and CUDNN using the pip wheels, since it is
+much easier! This method is only supported on x86_64, because NVIDIA has not
+released aarch64 CUDA pip packages.
 
 ```bash
 pip install --upgrade pip
 
-# NVIDIA CUDA 12 installation
+# CUDA 12 installation
 # Note: wheels only available on linux.
 pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
-If JAX detects the wrong version of the NVIDIA CUDA libraries, there are several things
-you need to check:
-
-* Make sure that `LD_LIBRARY_PATH` is not set, since `LD_LIBRARY_PATH` can
-  override the NVIDIA CUDA libraries.
-* Make sure that the NVIDIA CUDA libraries installed are those requested by JAX.
+If JAX detects the wrong version of the CUDA libraries, there are several things
+to check:
+* make sure that `LD_LIBRARY_PATH` is not set, since `LD_LIBRARY_PATH` can
+  override the CUDA libraries.
+* make sure that the CUDA libraries installed are those requested by JAX.
   Rerunning the installation command above should work.
 
-### pip installation: NVIDIA GPU (CUDA, installed locally, harder)
+### pip installation: GPU (CUDA, installed locally, harder)
 
-If you prefer to use a preinstalled copy of NVIDIA CUDA, you must first
-install NVIDIA [CUDA](https://developer.nvidia.com/cuda-downloads) and
-[cuDNN](https://developer.nvidia.com/CUDNN).
+If you prefer to use a preinstalled copy of CUDA, you must first
+install [CUDA](https://developer.nvidia.com/cuda-downloads) and
+[CuDNN](https://developer.nvidia.com/CUDNN).
 
 JAX provides pre-built CUDA-compatible wheels for **Linux x86_64 only**. Other
 combinations of operating system and architecture are possible, but require
-building from source (refer to {ref}`building-from-source` to learn more}.
+[building from source](developer.md#building-from-source).
 
 You should use an NVIDIA driver version that is at least as new as your
-[NVIDIA CUDA toolkit's corresponding driver version](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#cuda-major-component-versions__table-cuda-toolkit-driver-versions).
+[CUDA toolkit's corresponding driver version](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#cuda-major-component-versions__table-cuda-toolkit-driver-versions).
 If you need to use a newer CUDA toolkit with an older driver, for example
 on a cluster where you cannot update the NVIDIA driver easily, you may be
 able to use the
@@ -142,18 +126,19 @@ results.
 NCCL is an optional dependency, required only if you are performing multi-GPU
 computations.
 
-To install, run:
+To install, run
 
 ```bash
 pip install --upgrade pip
 
-# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 8.9 or newer.
+# Installs the wheel compatible with CUDA 12 and cuDNN 8.9 or newer.
 # Note: wheels only available on linux.
 pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
 ```
 
-**These `pip` installations do not work with Windows, and may fail silently; refer to the table
-[above](#supported-platforms).**
+**These `pip` installations do not work with Windows, and may fail silently; see
+[above](#installing-jax).**
 
 You can find your CUDA version with the command:
 
@@ -165,49 +150,43 @@ JAX uses `LD_LIBRARY_PATH` to find CUDA libraries and `PATH` to find binaries
 (`ptxas`, `nvlink`). Please make sure that these paths point to the correct CUDA
 installation.
 
-Please let the JAX team know on [the GitHub issue tracker](https://github.com/google/jax/issues)
-if you run into any errors or problems with the pre-built wheels.
+Please let us know on [the issue tracker](https://github.com/google/jax/issues)
+if you run into any errors or problems with the prebuilt wheels.
 
-(docker-containers-nvidia-gpu)=
-### NVIDIA GPU Docker containers
+### Docker containers: NVIDIA GPU
 
 NVIDIA provides the [JAX
 Toolbox](https://github.com/NVIDIA/JAX-Toolbox) containers, which are
 bleeding edge containers containing nightly releases of jax and some
 models/frameworks.
 
-## JAX nightly installation
+## Nightly installation
 
-Nightly releases reflect the state of the main JAX repository at the time they are
+Nightly releases reflect the state of the main repository at the time they are
 built, and may not pass the full test suite.
 
-- `jax`:
-
+* JAX:
 ```bash
 pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
 ```
 
-- `jaxlib` CPU:
-
+* Jaxlib CPU:
 ```bash
 pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 ```
 
-- `jaxlib` Google Cloud TPU:
-
+* Jaxlib TPU:
 ```bash
 pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-pip install -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install -U --pre libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 ```
 
-- `jaxlib` NVIDIA GPU (CUDA 12):
-
+* Jaxlib GPU (Cuda 12):
 ```bash
 pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html
 ```
 
-(install-google-tpu)=
-## Google Cloud TPU
+## Google TPU
 
 ### pip installation: Google Cloud TPU
 
@@ -215,53 +194,50 @@ JAX provides pre-built wheels for
 [Google Cloud TPU](https://cloud.google.com/tpu/docs/users-guide-tpu-vm).
 To install JAX along with appropriate versions of `jaxlib` and `libtpu`, you can run
 the following in your cloud TPU VM:
-
 ```bash
 pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 ```
 
-For users of Colab (https://colab.research.google.com/), be sure you are
-using *TPU v2* and not the older, deprecated TPU runtime.
+For interactive notebook users: Colab TPUs no longer support JAX as of
+JAX version 0.4. However, for an interactive TPU notebook in the cloud, you can
+use [Kaggle TPU notebooks](https://www.kaggle.com/docs/tpu), which fully
+support JAX.
 
-(install-apple-gpu)=
-## Apple Silicon GPU (ARM-based)
+## Apple GPU
 
-### pip installation: Apple ARM-based Silicon GPUs
+### pip installation: Apple GPUs
 
-Apple provides an experimental Metal plugin for Apple ARM-based GPU hardware. For details,
-refer to
+Apple provides an experimental Metal plugin for Apple GPU hardware. For details,
+see
 [Apple's JAX on Metal documentation](https://developer.apple.com/metal/jax/).
 
-**Note:** There are several caveats with the Metal plugin:
-
-* The Metal plugin is new and experimental and has a number of
+There are several caveats with the Metal plugin:
+* the Metal plugin is new and experimental and has a number of
   [known issues](https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+label%3A%22Apple+GPU+%28Metal%29+plugin%22).
   Please report any issues on the JAX issue tracker.
-* The Metal plugin currently requires very specific versions of `jax` and
+* the Metal plugin currently requires very specific versions of `jax` and
   `jaxlib`. This restriction will be relaxed over time as the plugin API
   matures.
 
-(install-amd-gpu)=
 ## AMD GPU
 
-JAX has experimental ROCm support. There are two ways to install JAX:
+JAX has experimental ROCM support. There are two ways to install JAX:
 
-* Use [AMD's Docker container](https://hub.docker.com/r/rocm/jax); or
-* Build from source (refer to {ref}`building-from-source` — a section called _Additional notes for building a ROCM `jaxlib` for AMD GPUs_).
+* use [AMD's docker container](https://hub.docker.com/r/rocm/jax), or
+* [build from source](developer.md#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus).
 
-## Conda (community-supported)
+## Conda
 
 ### Conda installation
 
-There is a community-supported Conda build of `jax`. To install it using `conda`,
-simply run:
+There is a community-supported Conda build of `jax`. To install using `conda`,
+simply run
 
 ```bash
 conda install jax -c conda-forge
 ```
 
-To install it on a machine with an NVIDIA GPU, run:
-
+To install on a machine with an NVIDIA GPU, run
 ```bash
 conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia
 ```
@@ -277,24 +253,20 @@ install the CUDA build on a machine without GPUs, follow the instructions in the
 [Tips & tricks](https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch)
 section of the `conda-forge` website.
 
-Go to the `conda-forge`
+See the `conda-forge`
 [jaxlib](https://github.com/conda-forge/jaxlib-feedstock#installing-jaxlib) and
 [jax](https://github.com/conda-forge/jax-feedstock#installing-jax) repositories
 for more details.
 
-
-(building-jax-from-source)=
 ## Building JAX from source
+See [Building JAX from source](developer.md#building-from-source).
 
-Refer to {ref}`building-from-source`.
+## Installing older jaxlib wheels
 
-## Installing older `jaxlib` wheels
-
-Due to storage limitations on the Python package index, the JAX team periodically removes
-older `jaxlib` wheels from the releases on http://pypi.org/project/jax. These can
-still be installed directly via the URLs here. For example:
-
-```bash
+Due to storage limitations on the Python package index, we periodically remove
+older jaxlib wheels from the releases on http://pypi.org/project/jax. These can
+still be installed directly via the URLs here; for example:
+```
 # Install jaxlib on CPU via the wheel archive
 pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
 
@@ -302,6 +274,6 @@ pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_
 pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
 ```
 For specific older GPU wheels, be sure to use the `jax_cuda_releases.html` URL; for example
-```bash
+```
 pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
diff --git a/docs/jax.numpy.rst b/docs/jax.numpy.rst
index 5ccf043d2..d866df52d 100644
--- a/docs/jax.numpy.rst
+++ b/docs/jax.numpy.rst
@@ -138,7 +138,6 @@ namespace; they are listed below.
     csingle
     cumprod
     cumsum
-    cumulative_sum
     deg2rad
     degrees
     delete
diff --git a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
index 54eb00d78..ebc2d55b3 100644
--- a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
+++ b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
@@ -412,8 +412,6 @@
     "id": "uRLpOcmNj_Vt"
    },
    "source": [
-    "The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.\n",
-    "\n",
     "By writing `PositionalSharding(ndarray_of_devices)`, we fix the device order and the initial shape. Then we can reshape it:"
    ]
   },
diff --git a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
index dad695d12..dca37d06b 100644
--- a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
+++ b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
@@ -192,8 +192,6 @@ sharding
 
 +++ {"id": "uRLpOcmNj_Vt"}
 
-The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.
-
 By writing `PositionalSharding(ndarray_of_devices)`, we fix the device order and the initial shape. Then we can reshape it:
 
 ```{code-cell}
diff --git a/docs/requirements.txt b/docs/requirements.txt
index e83ea4c92..104a7e06d 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,6 +1,6 @@
 absl-py
 ipython>=8.8.0  # 8.7.0 has ipython3 lexer error
-sphinx>=7.3.2  # 7.3.0 breaks sphinx-book-theme
+sphinx>=6.0.0
 sphinx-autodoc-typehints
 sphinx-book-theme>=1.0.1  # Older versions fail to pin pydata-sphinx-theme
 sphinx-copybutton>=0.5.0
diff --git a/docs/tutorials/automatic-differentiation.md b/docs/tutorials/automatic-differentiation.md
index 6175ee7da..c565a00db 100644
--- a/docs/tutorials/automatic-differentiation.md
+++ b/docs/tutorials/automatic-differentiation.md
@@ -162,7 +162,7 @@ Essentially, when using the `argnums` argument, if `f` is a Python function for
 (automatic-differentiation-nested-lists-tuples-and-dicts)=
 ## 3. Differentiating with respect to nested lists, tuples, and dicts
 
-Due to JAX's PyTree abstraction (see {ref}`pytrees-what-is-a-pytree`), differentiating with
+Due to JAX's PyTree abstraction (see {ref}`thinking-in-jax-pytrees`), differentiating with
 respect to standard Python containers just works, so use tuples, lists, and dicts (and arbitrary nesting) however you like.
 
 Continuing the previous example:
diff --git a/docs/tutorials/external-callbacks.md b/docs/tutorials/external-callbacks.md
index c30b373d4..9de1dea5d 100644
--- a/docs/tutorials/external-callbacks.md
+++ b/docs/tutorials/external-callbacks.md
@@ -42,7 +42,7 @@ def f(x):
 result = f(2)
 ```
 
-What is printed is not the runtime value, but the trace-time abstract value (if you're not familiar with *tracing* in JAX, a good primer can be found in {ref}`key-concepts-tracing`.
+What is printed is not the runtime value, but the trace-time abstract value (if you're not familiar with *tracing* in JAX, a good primer can be found in {ref}`thinking-in-jax`.
 
 To print the value at runtime, you need a callback, for example {func}`jax.debug.print` (you can learn more about debugging in {ref}`debugging`):
 
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
index 63754a1bb..0dadba6ce 100644
--- a/docs/tutorials/index.rst
+++ b/docs/tutorials/index.rst
@@ -19,13 +19,14 @@ JAX 101
 
    quickstart
    key-concepts
+   thinking-in-jax
    jit-compilation
    automatic-vectorization
    automatic-differentiation
    debugging
    random-numbers
    working-with-pytrees
-   sharded-computation
+   single-host-sharding
    stateful-computations
    simple-neural-network
 
@@ -53,3 +54,12 @@ JAX 301
    jax-primitives
    jaxpr
    advanced-compilation
+
+
+Reference
+---------
+
+.. toctree::
+   :maxdepth: 1
+
+   installation
diff --git a/docs/tutorials/installation.md b/docs/tutorials/installation.md
new file mode 100644
index 000000000..85879b53a
--- /dev/null
+++ b/docs/tutorials/installation.md
@@ -0,0 +1,311 @@
+(installation)=
+# Installing JAX
+
+This guide provides instructions for:
+
+- Installing JAX binary packages for supported platforms using `pip` or `conda`
+- Using Docker containers (for example {ref}`docker-containers-nvidia-gpu`)
+- {ref}`building-jax-from-source`
+
+**TL;DR** For most users, a typical JAX installation may look something like this:
+
+* **CPU-only (Linux/macOS/Windows)**
+  ```
+  pip install -U "jax[cpu]"
+  ```
+* **GPU (NVIDIA, CUDA 12, x86_64)**
+  ```
+  pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+  ```
+
+(install-supported-platforms)=
+## Supported platforms
+
+The table below shows all supported platforms and installation options. Check if your setup is supported; and if it says _"yes"_ or _"experimental"_, then click on the corresponding link to learn how to install JAX in greater detail.
+
+|                  | Linux, x86_64                        | Linux, aarch64      | macOS, Intel x86_64, AMD GPU   | macOS, Apple Silicon, ARM-based       | Windows, x86_64         | Windows WSL2, x86_64           |
+|------------------|---------------------------------------|--------------------------------|----------------------------------------|----------------------------------------|-------------------------|-----------------------------------------|
+| CPU              | {ref}`yes <install-cpu>`              | {ref}`yes <install-cpu>`        | {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>`|
+| NVIDIA GPU       | {ref}`yes <install-nvidia-gpu>`       | {ref}`yes <install-nvidia-gpu>` | no | n/a | no | {ref}`experimental <install-nvidia-gpu>` |
+| Google Cloud TPU | {ref}`yes <install-google-tpu>`       | n/a | n/a | n/a | n/a | n/a |
+| AMD GPU          | {ref}`experimental <install-amd-gpu>` | no | no | n/a | no | no |
+| Apple GPU    | n/a                                   | no | {ref}`experimental <install-apple-gpu>` | {ref}`experimental <install-apple-gpu>` | n/a |  n/a |
+
+
+(install-cpu)=
+## CPU
+
+### pip installation: CPU
+
+Currently, the JAX team releases `jaxlib` wheels for the following
+operating systems and architectures:
+
+- Linux, x86_64
+- macOS, Intel
+- macOS, Apple ARM-based
+- Windows, x86_64 (*experimental*)
+
+To install a CPU-only version of JAX, which might be useful for doing local
+development on a laptop, you can run:
+
+```bash
+pip install --upgrade pip
+pip install --upgrade "jax[cpu]"
+```
+
+On Windows, you may also need to install the
+[Microsoft Visual Studio 2019 Redistributable](https://learn.microsoft.com/en-US/cpp/windows/latest-supported-vc-redist?view=msvc-170#visual-studio-2015-2017-2019-and-2022)
+if it is not already installed on your machine.
+
+Other operating systems and architectures require building from source. Trying
+to pip install on other operating systems and architectures may lead to `jaxlib`
+not being installed alongside `jax`, although `jax` may successfully install
+(but fail at runtime).
+
+
+(install-nvidia-gpu)=
+## NVIDIA GPU
+
+JAX supports NVIDIA GPUs that have SM version 5.2 (Maxwell) or newer.
+Note that Kepler-series GPUs are no longer supported by JAX since
+NVIDIA has dropped support for Kepler GPUs in its software.
+
+You must first install the NVIDIA driver. You're
+recommended to install the newest driver available from NVIDIA, but the driver
+version must be >= 525.60.13 for CUDA 12 on Linux.
+
+If you need to use a newer CUDA toolkit with an older driver, for example
+on a cluster where you cannot update the NVIDIA driver easily, you may be
+able to use the
+[CUDA forward compatibility packages](https://docs.nvidia.com/deploy/cuda-compatibility/)
+that NVIDIA provides for this purpose.
+
+### pip installation: NVIDIA GPU (CUDA, installed via pip, easier)
+
+There are two ways to install JAX with NVIDIA GPU support:
+
+- Using NVIDIA CUDA and cuDNN installed from pip wheels
+- Using a self-installed CUDA/cuDNN
+
+The JAX team strongly recommends installing CUDA and cuDNN using the pip wheels,
+since it is much easier!
+
+This method is only supported on x86_64, because NVIDIA has not released aarch64
+CUDA pip packages.
+
+```bash
+pip install --upgrade pip
+
+# NVIDIA CUDA 12 installation
+# Note: wheels only available on linux.
+pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
+If JAX detects the wrong version of the NVIDIA CUDA libraries, there are several things
+you need to check:
+
+* Make sure that `LD_LIBRARY_PATH` is not set, since `LD_LIBRARY_PATH` can
+  override the NVIDIA CUDA libraries.
+* Make sure that the NVIDIA CUDA libraries installed are those requested by JAX.
+  Rerunning the installation command above should work.
+
+### pip installation: NVIDIA GPU (CUDA, installed locally, harder)
+
+If you prefer to use a preinstalled copy of NVIDIA CUDA, you must first
+install NVIDIA [CUDA](https://developer.nvidia.com/cuda-downloads) and
+[cuDNN](https://developer.nvidia.com/CUDNN).
+
+JAX provides pre-built CUDA-compatible wheels for **Linux x86_64 only**. Other
+combinations of operating system and architecture are possible, but require
+building from source (refer to {ref}`building-from-source` to learn more}.
+
+You should use an NVIDIA driver version that is at least as new as your
+[NVIDIA CUDA toolkit's corresponding driver version](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#cuda-major-component-versions__table-cuda-toolkit-driver-versions).
+If you need to use a newer CUDA toolkit with an older driver, for example
+on a cluster where you cannot update the NVIDIA driver easily, you may be
+able to use the
+[CUDA forward compatibility packages](https://docs.nvidia.com/deploy/cuda-compatibility/)
+that NVIDIA provides for this purpose.
+
+JAX currently ships one CUDA wheel variant:
+
+| Built with | Compatible with   |
+|------------|-------------------|
+| CUDA 12.3  | CUDA >=12.1       |
+| CUDNN 8.9  | CUDNN >=8.9, <9.0 |
+| NCCL 2.19  | NCCL >=2.18       |
+
+JAX checks the versions of your libraries, and will report an error if they are
+not sufficiently new.
+Setting the `JAX_SKIP_CUDA_CONSTRAINTS_CHECK` environment variable will disable
+the check, but using older versions of CUDA may lead to errors, or incorrect
+results.
+
+NCCL is an optional dependency, required only if you are performing multi-GPU
+computations.
+
+To install, run:
+
+```bash
+pip install --upgrade pip
+
+# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 8.9 or newer.
+# Note: wheels only available on linux.
+pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
+**These `pip` installations do not work with Windows, and may fail silently; refer to the table
+[above](#supported-platforms).**
+
+You can find your CUDA version with the command:
+
+```bash
+nvcc --version
+```
+
+JAX uses `LD_LIBRARY_PATH` to find CUDA libraries and `PATH` to find binaries
+(`ptxas`, `nvlink`). Please make sure that these paths point to the correct CUDA
+installation.
+
+Please let the JAX team know on [the GitHub issue tracker](https://github.com/google/jax/issues)
+if you run into any errors or problems with the pre-built wheels.
+
+(docker-containers-nvidia-gpu)=
+### NVIDIA GPU Docker containers
+
+NVIDIA provides the [JAX
+Toolbox](https://github.com/NVIDIA/JAX-Toolbox) containers, which are
+bleeding edge containers containing nightly releases of jax and some
+models/frameworks.
+
+## JAX nightly installation
+
+Nightly releases reflect the state of the main JAX repository at the time they are
+built, and may not pass the full test suite.
+
+- `jax`:
+
+```bash
+pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+```
+
+- `jaxlib` CPU:
+
+```bash
+pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+```
+
+- `jaxlib` Google Cloud TPU:
+
+```bash
+pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+pip install -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+```
+
+- `jaxlib` NVIDIA GPU (CUDA 12):
+
+```bash
+pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html
+```
+
+(install-google-tpu)=
+## Google Cloud TPU
+
+### pip installation: Google Cloud TPU
+
+JAX provides pre-built wheels for
+[Google Cloud TPU](https://cloud.google.com/tpu/docs/users-guide-tpu-vm).
+To install JAX along with appropriate versions of `jaxlib` and `libtpu`, you can run
+the following in your cloud TPU VM:
+
+```bash
+pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+```
+
+For interactive notebook users: Colab TPUs no longer support JAX as of
+JAX version 0.4. However, for an interactive TPU notebook in the cloud, you can
+use [Kaggle TPU notebooks](https://www.kaggle.com/docs/tpu), which fully
+support JAX.
+
+(install-apple-gpu)=
+## Apple Silicon GPU (ARM-based)
+
+### pip installation: Apple ARM-based Silicon GPUs
+
+Apple provides an experimental Metal plugin for Apple ARM-based GPU hardware. For details,
+refer to
+[Apple's JAX on Metal documentation](https://developer.apple.com/metal/jax/).
+
+**Note:** There are several caveats with the Metal plugin:
+
+* The Metal plugin is new and experimental and has a number of
+  [known issues](https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+label%3A%22Apple+GPU+%28Metal%29+plugin%22).
+  Please report any issues on the JAX issue tracker.
+* The Metal plugin currently requires very specific versions of `jax` and
+  `jaxlib`. This restriction will be relaxed over time as the plugin API
+  matures.
+
+(install-amd-gpu)=
+## AMD GPU
+
+JAX has experimental ROCm support. There are two ways to install JAX:
+
+* Use [AMD's Docker container](https://hub.docker.com/r/rocm/jax); or
+* Build from source (refer to {ref}`building-from-source` — a section called _Additional notes for building a ROCM `jaxlib` for AMD GPUs_).
+
+## Conda (community-supported)
+
+### Conda installation
+
+There is a community-supported Conda build of `jax`. To install it using `conda`,
+simply run:
+
+```bash
+conda install jax -c conda-forge
+```
+
+To install it on a machine with an NVIDIA GPU, run:
+
+```bash
+conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia
+```
+
+Note the `cudatoolkit` distributed by `conda-forge` is missing `ptxas`, which
+JAX requires. You must therefore either install the `cuda-nvcc` package from
+the `nvidia` channel, or install CUDA on your machine separately so that `ptxas`
+is in your path. The channel order above is important (`conda-forge` before
+`nvidia`).
+
+If you would like to override which release of CUDA is used by JAX, or to
+install the CUDA build on a machine without GPUs, follow the instructions in the
+[Tips & tricks](https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch)
+section of the `conda-forge` website.
+
+Go to the `conda-forge`
+[jaxlib](https://github.com/conda-forge/jaxlib-feedstock#installing-jaxlib) and
+[jax](https://github.com/conda-forge/jax-feedstock#installing-jax) repositories
+for more details.
+
+
+(building-jax-from-source)=
+## Building JAX from source
+
+Refer to {ref}`building-from-source`.
+
+## Installing older `jaxlib` wheels
+
+Due to storage limitations on the Python package index, the JAX team periodically removes
+older `jaxlib` wheels from the releases on http://pypi.org/project/jax. These can
+still be installed directly via the URLs here. For example:
+
+```bash
+# Install jaxlib on CPU via the wheel archive
+pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
+
+# Install the jaxlib 0.3.25 CPU wheel directly
+pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
+```
+For specific older GPU wheels, be sure to use the `jax_cuda_releases.html` URL; for example
+```bash
+pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
diff --git a/docs/tutorials/key-concepts.md b/docs/tutorials/key-concepts.md
index a6dbf5f4b..46a9d2669 100644
--- a/docs/tutorials/key-concepts.md
+++ b/docs/tutorials/key-concepts.md
@@ -58,7 +58,7 @@ x.sharding
 
 Here the array is on a single device, but in general a JAX array can be
 sharded across multiple devices, or even multiple hosts.
-To read more about sharded arrays and parallel computation, refer to {ref}`sharded-computation`
+To read more about sharded arrays and parallel computation, refer to {ref}`single-host-sharding`
 
 (key-concepts-transformations)=
 ## Transformations
diff --git a/docs/tutorials/sharded-computation.ipynb b/docs/tutorials/sharded-computation.ipynb
deleted file mode 100644
index 55f19716f..000000000
--- a/docs/tutorials/sharded-computation.ipynb
+++ /dev/null
@@ -1,764 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "(sharded-computation)=\n",
-    "# Introduction to sharded computation\n",
-    "\n",
-    "JAX's {class}`jax.Array` object is designed with distributed data and computation in mind.\n",
-    "\n",
-    "This section will cover three modes of parallel computation:\n",
-    "\n",
-    "- Automatic parallelism via {func}`jax.jit`, in which we let the compiler choose the optimal computation strategy\n",
-    "- Semi-automatic parallelism using {func}`jax.jit` and {func}`jax.lax.with_sharding_constraint`\n",
-    "- Fully manual parallelism using {func}`jax.experimental.shard_map.shard_map`\n",
-    "\n",
-    "These examples will be run on Colab's free TPU runtime, which provides eight devices to work with:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {
-    "outputId": "18905ae4-7b5e-4bb9-acb4-d8ab914cb456"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
-       " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
-       " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
-       " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
-       " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
-       " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
-       " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
-       " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
-      ]
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import jax\n",
-    "jax.devices()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Key concept: data sharding\n",
-    "\n",
-    "Key to all of the distributed computation approaches below is the concept of *data sharding*, which describes how data is laid out on the available devices.\n",
-    "\n",
-    "Each concrete {class}`jax.Array` object has a `sharding` attribute and a `devices()` method that can give you insight into how the underlying data are stored. In the simplest cases, arrays are sharded on a single device:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "outputId": "39fdbb79-d5c0-4ea6-8b20-88b2c502a27a"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)}"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import jax.numpy as jnp\n",
-    "arr = jnp.arange(32.0).reshape(4, 8)\n",
-    "arr.devices()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "outputId": "536f773a-7ef4-4526-c58b-ab4d486bf5a1"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "SingleDeviceSharding(device=TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0))"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "arr.sharding"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "For a more visual representation of the storage layout, the {mod}`jax.debug` module provides some helpers to visualize the sharding of an array:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "outputId": "74a793e9-b13b-4d07-d8ec-7e25c547036d"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                      TPU 0                       </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                      \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                       \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "jax.debug.visualize_array_sharding(arr)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To create an array with a non-trivial sharding, we can define a `sharding` specification for the array and pass this to {func}`jax.device_put`.\n",
-    "Here we'll define a {class}`~jax.sharding.NamedSharding`, which specifies an N-dimensional grid of devices with named axes:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "outputId": "0b397dba-3ddc-4aca-f002-2beab7e6b8a5"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "NamedSharding(mesh=Mesh('x': 2, 'y': 4), spec=PartitionSpec('x', 'y'))\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Pardon the boilerplate; constructing a sharding will become easier soon!\n",
-    "from jax.sharding import Mesh\n",
-    "from jax.sharding import PartitionSpec\n",
-    "from jax.sharding import NamedSharding\n",
-    "from jax.experimental import mesh_utils\n",
-    "\n",
-    "P = jax.sharding.PartitionSpec\n",
-    "devices = mesh_utils.create_device_mesh((2, 4))\n",
-    "mesh = jax.sharding.Mesh(devices, ('x', 'y'))\n",
-    "sharding = jax.sharding.NamedSharding(mesh, P('x', 'y'))\n",
-    "print(sharding)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Passing this `sharding` to {func}`jax.device_put`, we obtain a sharded array:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "outputId": "c8ceedba-05ca-4156-e6e4-1e98bb664a66"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[[ 0.  1.  2.  3.  4.  5.  6.  7.]\n",
-      " [ 8.  9. 10. 11. 12. 13. 14. 15.]\n",
-      " [16. 17. 18. 19. 20. 21. 22. 23.]\n",
-      " [24. 25. 26. 27. 28. 29. 30. 31.]]\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   TPU 0    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">   TPU 1    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">   TPU 2    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">   TPU 3    </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">   TPU 6    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">   TPU 7    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">   TPU 4    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">   TPU 5    </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m   \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m    \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m   \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m    \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m   \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m    \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m   \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mTPU 6\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m    \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m   \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mTPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m    \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m   \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mTPU 4\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m    \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m   \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mTPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m    \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "arr_sharded = jax.device_put(arr, sharding)\n",
-    "\n",
-    "print(arr_sharded)\n",
-    "jax.debug.visualize_array_sharding(arr_sharded)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.\n",
-    "\n",
-    "\n",
-    "\n",
-    "## Automatic parallelism via `jit`\n",
-    "Once you have sharded data, the easiest way to do parallel computation is to simply pass the data to a JIT-compiled function!\n",
-    "The XLA compiler behind `jit` includes heuristics for optimizing computations across multiple devices.\n",
-    "In the simplest of cases, those heuristics boil down to *computation follows data*.\n",
-    "\n",
-    "For example, here's a simple element-wise function: the computation for each shard will be performed on the device associated with that shard, and the output is sharded in the same way:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "outputId": "de46f86a-6907-49c8-f36c-ed835e78bc3d"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "shardings match: True\n"
-     ]
-    }
-   ],
-   "source": [
-    "@jax.jit\n",
-    "def f_elementwise(x):\n",
-    "  return 2 * jnp.sin(x) + 1\n",
-    "\n",
-    "result = f_elementwise(arr_sharded)\n",
-    "\n",
-    "print(\"shardings match:\", result.sharding == arr_sharded.sharding)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "As computations get more complex, the compiler makes decisions about how to best propagate the sharding of the data.\n",
-    "Here we sum along the leading axis of `x`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {
-    "outputId": "90c3b997-3653-4a7b-c8ff-12a270f11d02"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> TPU 0,6 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> TPU 1,7 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> TPU 2,4 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> TPU 3,5 </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,6\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1,7\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2,4\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3,5\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[48. 52. 56. 60. 64. 68. 72. 76.]\n"
-     ]
-    }
-   ],
-   "source": [
-    "@jax.jit\n",
-    "def f_contract(x):\n",
-    "  return x.sum(axis=0)\n",
-    "\n",
-    "result = f_contract(arr_sharded)\n",
-    "jax.debug.visualize_array_sharding(result)\n",
-    "print(result)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The result is partially replicated: that is, the first two elements of the array are replicated on devices `0` and `6`, the second on `1` and `7`, and so on.\n",
-    "\n",
-    "\n",
-    "\n",
-    "## Semi-automated sharding with constraints\n",
-    "\n",
-    "If you'd like to have some control over the sharding used within a particular computation, JAX offers the {func}`~jax.lax.with_sharding_constraint` function.\n",
-    "\n",
-    "For example, suppose that within `f_contract` above, you'd prefer the output not to be partially-replicated, but rather to be fully sharded across the eight devices:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {
-    "outputId": "8468f5c6-76ca-4367-c9f2-93c723687cfd"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  TPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  TPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 3  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  TPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  TPU 7  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  TPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  TPU 5  </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mTPU 6\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mTPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mTPU 4\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mTPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[48. 52. 56. 60. 64. 68. 72. 76.]\n"
-     ]
-    }
-   ],
-   "source": [
-    "@jax.jit\n",
-    "def f_contract_2(x):\n",
-    "  out = x.sum(axis=0)\n",
-    "  # mesh = jax.create_mesh((8,), 'x')\n",
-    "  devices = mesh_utils.create_device_mesh(8)\n",
-    "  mesh = jax.sharding.Mesh(devices, 'x')\n",
-    "  sharding = jax.sharding.NamedSharding(mesh, P('x'))\n",
-    "  return jax.lax.with_sharding_constraint(out, sharding)\n",
-    "\n",
-    "result = f_contract_2(arr_sharded)\n",
-    "jax.debug.visualize_array_sharding(result)\n",
-    "print(result)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This gives you a function with the particular output sharding you'd like.\n",
-    "\n",
-    "\n",
-    "\n",
-    "## Manual parallelism with `shard_map`\n",
-    "\n",
-    "In the automatic parallelism methods explored above, you can write a function as if you're operating on the full dataset, and `jit` will split that computation across multiple devices.\n",
-    "By contrast, with `shard_map` you write the function that will handle a single shard of data, and `shard_map` will construct the full function.\n",
-    "\n",
-    "`shard_map` works by mapping a function across a particular *mesh* of devices:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {
-    "outputId": "435c32f3-557a-4676-c11b-17e6bab8c1e2"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([ 1.        ,  2.682942  ,  2.818595  ,  1.28224   , -0.513605  ,\n",
-       "       -0.9178486 ,  0.44116896,  2.3139732 ,  2.9787164 ,  1.824237  ,\n",
-       "       -0.08804226, -0.99998045, -0.07314599,  1.8403342 ,  2.9812148 ,\n",
-       "        2.3005757 ,  0.42419332, -0.92279506, -0.50197446,  1.2997544 ,\n",
-       "        2.8258905 ,  2.6733112 ,  0.98229736, -0.69244075, -0.81115675,\n",
-       "        0.7352965 ,  2.525117  ,  2.912752  ,  1.5418116 , -0.32726777,\n",
-       "       -0.97606325,  0.19192469], dtype=float32)"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from jax.experimental.shard_map import shard_map\n",
-    "P = jax.sharding.PartitionSpec\n",
-    "mesh = jax.sharding.Mesh(jax.devices(), 'x')\n",
-    "\n",
-    "f_elementwise_sharded = shard_map(\n",
-    "    f_elementwise,\n",
-    "    mesh=mesh,\n",
-    "    in_specs=P('x'),\n",
-    "    out_specs=P('x'))\n",
-    "\n",
-    "arr = jnp.arange(32)\n",
-    "f_elementwise_sharded(arr)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The function you write only \"sees\" a single batch of the data, which we can see by printing the device local shape:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {
-    "outputId": "99a3dc6e-154a-4ef6-8eaa-3dd0b68fb1da"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "global shape: x.shape=(32,)\n",
-      "device local shape: x.shape=(4,)\n"
-     ]
-    }
-   ],
-   "source": [
-    "x = jnp.arange(32)\n",
-    "print(f\"global shape: {x.shape=}\")\n",
-    "\n",
-    "def f(x):\n",
-    "  print(f\"device local shape: {x.shape=}\")\n",
-    "  return x * 2\n",
-    "\n",
-    "y = shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P('x'))(x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Because each of your functions only sees the device-local part of the data, it means that aggregation-like functions require some extra thought.\n",
-    "For example, here's what a `shard_map` of a `sum` looks like:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {
-    "outputId": "1e9a45f5-5418-4246-c75b-f9bc6dcbbe72"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([  6,  22,  38,  54,  70,  86, 102, 118], dtype=int32)"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "def f(x):\n",
-    "  return jnp.sum(x, keepdims=True)\n",
-    "\n",
-    "shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P('x'))(x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Our function `f` operates separately on each shard, and the resulting summation reflects this.\n",
-    "If we want to sum across shards, we need to explicitly request it using collective operations like {func}`jax.lax.psum`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {
-    "outputId": "4fd29e80-4fee-42b7-ff80-29f9887ab38d"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array(496, dtype=int32)"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "def f(x):\n",
-    "  sum_in_shard = x.sum()\n",
-    "  return jax.lax.psum(sum_in_shard, 'x')\n",
-    "\n",
-    "shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P())(x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Because the output no longer has a sharded dimension, we set `out_specs=P()`.\n",
-    "\n",
-    "\n",
-    "\n",
-    "## Comparing the three approaches\n",
-    "\n",
-    "With these concepts fresh in our mind, let's compare the three approaches for a simple neural network layer.\n",
-    "We'll define our canonical function like this:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {
-    "id": "1TdhfTsoiqS1"
-   },
-   "outputs": [],
-   "source": [
-    "@jax.jit\n",
-    "def layer(x, weights, bias):\n",
-    "  return jax.nn.sigmoid(x @ weights + bias)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {
-    "outputId": "f3007fe4-f6f3-454e-e7c5-3638de484c0a"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([0.02138912, 0.893112  , 0.59892005, 0.97742504], dtype=float32)"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import numpy as np\n",
-    "rng = np.random.default_rng(0)\n",
-    "\n",
-    "x = rng.normal(size=(32,))\n",
-    "weights = rng.normal(size=(32, 4))\n",
-    "bias = rng.normal(size=(4,))\n",
-    "\n",
-    "layer(x, weights, bias)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can automatically run this in a distributed manner using {func}`jax.jit` and passing appropriately sharded data.\n",
-    "If we shard the leading axis of both `x` and `weights` in the same way, then the matrix multiplication will autoatically happen in parallel:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {
-    "outputId": "80be899e-8dbc-4bfc-acd2-0f3d554a0aa5"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([0.02138912, 0.893112  , 0.59892005, 0.97742504], dtype=float32)"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "P = jax.sharding.PartitionSpec\n",
-    "mesh = jax.sharding.Mesh(jax.devices(), 'x')\n",
-    "sharding = jax.sharding.NamedSharding(mesh, P('x'))\n",
-    "\n",
-    "x_sharded = jax.device_put(x, sharding)\n",
-    "weights_sharded = jax.device_put(weights, sharding)\n",
-    "\n",
-    "layer(x_sharded, weights_sharded, bias)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Alternatively, we can use {func}`jax.lax.with_sharding_constraint` in the function to automatically distribute unsharded inputs:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {
-    "outputId": "bb63e8da-ff4f-4e95-f083-10584882daf4"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([0.02138914, 0.89311206, 0.5989201 , 0.97742516], dtype=float32)"
-      ]
-     },
-     "execution_count": 17,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "@jax.jit\n",
-    "def layer_auto(x, weights, bias):\n",
-    "  x = jax.lax.with_sharding_constraint(x, sharding)\n",
-    "  weights = jax.lax.with_sharding_constraint(weights, sharding)\n",
-    "  return layer(x, weights, bias)\n",
-    "\n",
-    "layer_auto(x, weights, bias)  # pass in unsharded inputs"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Finally, we can do the same thing with `shard_map`, using `psum` to indicate the cross-shard collective required for the matrix product:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {
-    "outputId": "568d1c85-39a7-4dba-f09a-0e4f7c2ea918"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([0.02138914, 0.89311206, 0.5989201 , 0.97742516], dtype=float32)"
-      ]
-     },
-     "execution_count": 18,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from functools import partial\n",
-    "\n",
-    "@jax.jit\n",
-    "@partial(shard_map, mesh=mesh,\n",
-    "         in_specs=(P('x'), P('x', None), P(None)),\n",
-    "         out_specs=P(None))\n",
-    "def layer_sharded(x, weights, bias):\n",
-    "  return jax.nn.sigmoid(jax.lax.psum(x @ weights, 'x') + bias)\n",
-    "\n",
-    "layer_sharded(x, weights, bias)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This section has been a brief introduction of sharded and parallel computation;\n",
-    "for more discussion of `shard_map`, see {doc}`../notebooks/shard_map`."
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "TPU",
-  "colab": {
-   "gpuType": "V28",
-   "provenance": [],
-   "toc_visible": true
-  },
-  "jupytext": {
-   "formats": "ipynb,md:myst"
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/docs/tutorials/sharded-computation.md b/docs/tutorials/sharded-computation.md
deleted file mode 100644
index aac37d0fb..000000000
--- a/docs/tutorials/sharded-computation.md
+++ /dev/null
@@ -1,303 +0,0 @@
----
-jupytext:
-  formats: ipynb,md:myst
-  text_representation:
-    extension: .md
-    format_name: myst
-    format_version: 0.13
-    jupytext_version: 1.16.1
-kernelspec:
-  display_name: Python 3
-  name: python3
----
-
-(sharded-computation)=
-# Introduction to sharded computation
-
-JAX's {class}`jax.Array` object is designed with distributed data and computation in mind.
-
-This section will cover three modes of parallel computation:
-
-- Automatic parallelism via {func}`jax.jit`, in which we let the compiler choose the optimal computation strategy
-- Semi-automatic parallelism using {func}`jax.jit` and {func}`jax.lax.with_sharding_constraint`
-- Fully manual parallelism using {func}`jax.experimental.shard_map.shard_map`
-
-These examples will be run on Colab's free TPU runtime, which provides eight devices to work with:
-
-```{code-cell}
-:outputId: 18905ae4-7b5e-4bb9-acb4-d8ab914cb456
-
-import jax
-jax.devices()
-```
-
-## Key concept: data sharding
-
-Key to all of the distributed computation approaches below is the concept of *data sharding*, which describes how data is laid out on the available devices.
-
-Each concrete {class}`jax.Array` object has a `sharding` attribute and a `devices()` method that can give you insight into how the underlying data are stored. In the simplest cases, arrays are sharded on a single device:
-
-```{code-cell}
-:outputId: 39fdbb79-d5c0-4ea6-8b20-88b2c502a27a
-
-import jax.numpy as jnp
-arr = jnp.arange(32.0).reshape(4, 8)
-arr.devices()
-```
-
-```{code-cell}
-:outputId: 536f773a-7ef4-4526-c58b-ab4d486bf5a1
-
-arr.sharding
-```
-
-For a more visual representation of the storage layout, the {mod}`jax.debug` module provides some helpers to visualize the sharding of an array:
-
-```{code-cell}
-:outputId: 74a793e9-b13b-4d07-d8ec-7e25c547036d
-
-jax.debug.visualize_array_sharding(arr)
-```
-
-To create an array with a non-trivial sharding, we can define a `sharding` specification for the array and pass this to {func}`jax.device_put`.
-Here we'll define a {class}`~jax.sharding.NamedSharding`, which specifies an N-dimensional grid of devices with named axes:
-
-```{code-cell}
-:outputId: 0b397dba-3ddc-4aca-f002-2beab7e6b8a5
-
-# Pardon the boilerplate; constructing a sharding will become easier soon!
-from jax.sharding import Mesh
-from jax.sharding import PartitionSpec
-from jax.sharding import NamedSharding
-from jax.experimental import mesh_utils
-
-P = jax.sharding.PartitionSpec
-devices = mesh_utils.create_device_mesh((2, 4))
-mesh = jax.sharding.Mesh(devices, ('x', 'y'))
-sharding = jax.sharding.NamedSharding(mesh, P('x', 'y'))
-print(sharding)
-```
-
-Passing this `sharding` to {func}`jax.device_put`, we obtain a sharded array:
-
-```{code-cell}
-:outputId: c8ceedba-05ca-4156-e6e4-1e98bb664a66
-
-arr_sharded = jax.device_put(arr, sharding)
-
-print(arr_sharded)
-jax.debug.visualize_array_sharding(arr_sharded)
-```
-
-The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.
-
-
-
-## Automatic parallelism via `jit`
-Once you have sharded data, the easiest way to do parallel computation is to simply pass the data to a JIT-compiled function!
-The XLA compiler behind `jit` includes heuristics for optimizing computations across multiple devices.
-In the simplest of cases, those heuristics boil down to *computation follows data*.
-
-For example, here's a simple element-wise function: the computation for each shard will be performed on the device associated with that shard, and the output is sharded in the same way:
-
-```{code-cell}
-:outputId: de46f86a-6907-49c8-f36c-ed835e78bc3d
-
-@jax.jit
-def f_elementwise(x):
-  return 2 * jnp.sin(x) + 1
-
-result = f_elementwise(arr_sharded)
-
-print("shardings match:", result.sharding == arr_sharded.sharding)
-```
-
-As computations get more complex, the compiler makes decisions about how to best propagate the sharding of the data.
-Here we sum along the leading axis of `x`:
-
-```{code-cell}
-:outputId: 90c3b997-3653-4a7b-c8ff-12a270f11d02
-
-@jax.jit
-def f_contract(x):
-  return x.sum(axis=0)
-
-result = f_contract(arr_sharded)
-jax.debug.visualize_array_sharding(result)
-print(result)
-```
-
-The result is partially replicated: that is, the first two elements of the array are replicated on devices `0` and `6`, the second on `1` and `7`, and so on.
-
-
-
-## Semi-automated sharding with constraints
-
-If you'd like to have some control over the sharding used within a particular computation, JAX offers the {func}`~jax.lax.with_sharding_constraint` function.
-
-For example, suppose that within `f_contract` above, you'd prefer the output not to be partially-replicated, but rather to be fully sharded across the eight devices:
-
-```{code-cell}
-:outputId: 8468f5c6-76ca-4367-c9f2-93c723687cfd
-
-@jax.jit
-def f_contract_2(x):
-  out = x.sum(axis=0)
-  # mesh = jax.create_mesh((8,), 'x')
-  devices = mesh_utils.create_device_mesh(8)
-  mesh = jax.sharding.Mesh(devices, 'x')
-  sharding = jax.sharding.NamedSharding(mesh, P('x'))
-  return jax.lax.with_sharding_constraint(out, sharding)
-
-result = f_contract_2(arr_sharded)
-jax.debug.visualize_array_sharding(result)
-print(result)
-```
-
-This gives you a function with the particular output sharding you'd like.
-
-
-
-## Manual parallelism with `shard_map`
-
-In the automatic parallelism methods explored above, you can write a function as if you're operating on the full dataset, and `jit` will split that computation across multiple devices.
-By contrast, with `shard_map` you write the function that will handle a single shard of data, and `shard_map` will construct the full function.
-
-`shard_map` works by mapping a function across a particular *mesh* of devices:
-
-```{code-cell}
-:outputId: 435c32f3-557a-4676-c11b-17e6bab8c1e2
-
-from jax.experimental.shard_map import shard_map
-P = jax.sharding.PartitionSpec
-mesh = jax.sharding.Mesh(jax.devices(), 'x')
-
-f_elementwise_sharded = shard_map(
-    f_elementwise,
-    mesh=mesh,
-    in_specs=P('x'),
-    out_specs=P('x'))
-
-arr = jnp.arange(32)
-f_elementwise_sharded(arr)
-```
-
-The function you write only "sees" a single batch of the data, which we can see by printing the device local shape:
-
-```{code-cell}
-:outputId: 99a3dc6e-154a-4ef6-8eaa-3dd0b68fb1da
-
-x = jnp.arange(32)
-print(f"global shape: {x.shape=}")
-
-def f(x):
-  print(f"device local shape: {x.shape=}")
-  return x * 2
-
-y = shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P('x'))(x)
-```
-
-Because each of your functions only sees the device-local part of the data, it means that aggregation-like functions require some extra thought.
-For example, here's what a `shard_map` of a `sum` looks like:
-
-```{code-cell}
-:outputId: 1e9a45f5-5418-4246-c75b-f9bc6dcbbe72
-
-def f(x):
-  return jnp.sum(x, keepdims=True)
-
-shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P('x'))(x)
-```
-
-Our function `f` operates separately on each shard, and the resulting summation reflects this.
-If we want to sum across shards, we need to explicitly request it using collective operations like {func}`jax.lax.psum`:
-
-```{code-cell}
-:outputId: 4fd29e80-4fee-42b7-ff80-29f9887ab38d
-
-def f(x):
-  sum_in_shard = x.sum()
-  return jax.lax.psum(sum_in_shard, 'x')
-
-shard_map(f, mesh=mesh, in_specs=P('x'), out_specs=P())(x)
-```
-
-Because the output no longer has a sharded dimension, we set `out_specs=P()`.
-
-
-
-## Comparing the three approaches
-
-With these concepts fresh in our mind, let's compare the three approaches for a simple neural network layer.
-We'll define our canonical function like this:
-
-```{code-cell}
-:id: 1TdhfTsoiqS1
-
-@jax.jit
-def layer(x, weights, bias):
-  return jax.nn.sigmoid(x @ weights + bias)
-```
-
-```{code-cell}
-:outputId: f3007fe4-f6f3-454e-e7c5-3638de484c0a
-
-import numpy as np
-rng = np.random.default_rng(0)
-
-x = rng.normal(size=(32,))
-weights = rng.normal(size=(32, 4))
-bias = rng.normal(size=(4,))
-
-layer(x, weights, bias)
-```
-
-We can automatically run this in a distributed manner using {func}`jax.jit` and passing appropriately sharded data.
-If we shard the leading axis of both `x` and `weights` in the same way, then the matrix multiplication will autoatically happen in parallel:
-
-```{code-cell}
-:outputId: 80be899e-8dbc-4bfc-acd2-0f3d554a0aa5
-
-P = jax.sharding.PartitionSpec
-mesh = jax.sharding.Mesh(jax.devices(), 'x')
-sharding = jax.sharding.NamedSharding(mesh, P('x'))
-
-x_sharded = jax.device_put(x, sharding)
-weights_sharded = jax.device_put(weights, sharding)
-
-layer(x_sharded, weights_sharded, bias)
-```
-
-Alternatively, we can use {func}`jax.lax.with_sharding_constraint` in the function to automatically distribute unsharded inputs:
-
-```{code-cell}
-:outputId: bb63e8da-ff4f-4e95-f083-10584882daf4
-
-@jax.jit
-def layer_auto(x, weights, bias):
-  x = jax.lax.with_sharding_constraint(x, sharding)
-  weights = jax.lax.with_sharding_constraint(weights, sharding)
-  return layer(x, weights, bias)
-
-layer_auto(x, weights, bias)  # pass in unsharded inputs
-```
-
-Finally, we can do the same thing with `shard_map`, using `psum` to indicate the cross-shard collective required for the matrix product:
-
-```{code-cell}
-:outputId: 568d1c85-39a7-4dba-f09a-0e4f7c2ea918
-
-from functools import partial
-
-@jax.jit
-@partial(shard_map, mesh=mesh,
-         in_specs=(P('x'), P('x', None), P(None)),
-         out_specs=P(None))
-def layer_sharded(x, weights, bias):
-  return jax.nn.sigmoid(jax.lax.psum(x @ weights, 'x') + bias)
-
-layer_sharded(x, weights, bias)
-```
-
-This section has been a brief introduction of sharded and parallel computation;
-for more discussion of `shard_map`, see {doc}`../notebooks/shard_map`.
diff --git a/docs/tutorials/single-host-sharding.md b/docs/tutorials/single-host-sharding.md
new file mode 100644
index 000000000..c21a90b19
--- /dev/null
+++ b/docs/tutorials/single-host-sharding.md
@@ -0,0 +1,6 @@
+(single-host-sharding)=
+# Sharded data on a single host
+
+```{note}
+This is a placeholder for a section in the new {ref}`jax-tutorials`.
+```
diff --git a/docs/tutorials/stateful-computations.md b/docs/tutorials/stateful-computations.md
index b802be0e0..318d4c4c6 100644
--- a/docs/tutorials/stateful-computations.md
+++ b/docs/tutorials/stateful-computations.md
@@ -1,235 +1,8 @@
----
-jupytext:
-  formats: md:myst
-  text_representation:
-    extension: .md
-    format_name: myst
-    format_version: 0.13
-    jupytext_version: 1.16.1
-kernelspec:
-  display_name: Python 3
-  language: python
-  name: python3
----
+# Stateful computations
 
-# Stateful Computations
+```{note}
+This is a placeholder for a section in the new {ref}`jax-tutorials`.
 
-JAX transformations like {func}`~jax.jit`, {func}`~jax.vmap`, {func}`~jax.grad`, require the functions
-they wrap to be pure: that is, functions whose outputs depend *solely* on the inputs, and which have
-no side effects such as updating of global state.
-You can find a discussion of this in [JAX sharp bits: Pure functions](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions).
-
-This constraint can pose some challenges in the context of machine learning, where state may exist in
-many forms. For example:
-
-* model parameters,
-* optimizer state, and
-* stateful layers, such as [BatchNorm](https://en.wikipedia.org/wiki/Batch_normalization).
-
-This section offers some advice of how to properly handle state in a JAX program.
-
-## A simple example: Counter
-
-Let's start by looking at a simple stateful program: a counter.
-
-```{code-cell}
-import jax
-import jax.numpy as jnp
-
-class Counter:
-  """A simple counter."""
-
-  def __init__(self):
-    self.n = 0
-
-  def count(self) -> int:
-    """Increments the counter and returns the new value."""
-    self.n += 1
-    return self.n
-
-  def reset(self):
-    """Resets the counter to zero."""
-    self.n = 0
-
-
-counter = Counter()
-
-for _ in range(3):
-  print(counter.count())
-```
-
-The counter's `n` attribute maintains the counter's _state_ between successive calls of `count`. It is modified as a side effect of calling `count`.
-
-Let's say we want to count fast, so we JIT-compile the `count` method.
-(In this example, this wouldn't actually help speed anyway, for many reasons, but treat this as a toy model of JIT-compiling the update of model parameters, where {func}`~jax.jit` makes an enormous difference).
-
-```{code-cell}
-counter.reset()
-fast_count = jax.jit(counter.count)
-
-for _ in range(3):
-  print(fast_count())
-```
-
-Oh no! Our counter isn't working. This is because the line
-```
-self.n += 1
-```
-in `count` involves a side effect: it modifies the input counter in-place, and so this function is not supported by `jit`.
-Such side effects are executed only once when the function is first traced, and subsequent calls will not repeat the side effect.
-So, how do we fix it?
-
-## The solution: explicit state
-
-Part of the problem with our counter was that the returned value didn't depend on the arguments, meaning a constant was "baked into" the compiled output. But it shouldn't be a constant -- it should depend on the state. Well, then why don't we make the state into an argument?
-
-```{code-cell}
-CounterState = int
-
-class CounterV2:
-
-  def count(self, n: CounterState) -> tuple[int, CounterState]:
-    # You could just return n+1, but here we separate its role as 
-    # the output and as the counter state for didactic purposes.
-    return n+1, n+1
-
-  def reset(self) -> CounterState:
-    return 0
-
-counter = CounterV2()
-state = counter.reset()
-
-for _ in range(3):
-  value, state = counter.count(state)
-  print(value)
-```
-
-In this new version of `Counter`, we moved `n` to be an argument of `count`, and added another return value that represents the new, updated, state. To use this counter, we now need to keep track of the state explicitly. But in return, we can now safely `jax.jit` this counter:
-
-```{code-cell}
-state = counter.reset()
-fast_count = jax.jit(counter.count)
-
-for _ in range(3):
-  value, state = fast_count(state)
-  print(value)
-```
-
-## A general strategy
-
-We can apply the same process to any stateful method to convert it into a stateless one. We took a class of the form
-
-```python
-class StatefulClass
-
-  state: State
-
-  def stateful_method(*args, **kwargs) -> Output:
+For the time being, you may find some related content in the old documentation:
+- {doc}`../jax-101/07-state`
 ```
-
-and turned it into a class of the form
-
-```python
-class StatelessClass
-
-  def stateless_method(state: State, *args, **kwargs) -> (Output, State):
-```
-
-This is a common [functional programming](https://en.wikipedia.org/wiki/Functional_programming) pattern, and, essentially, is the way that state is handled in all JAX programs.
-
-Notice that the need for a class becomes less clear once we have rewritten it this way. We could just keep `stateless_method`, since the class is no longer doing any work.
-This is because, like the strategy we just applied, object-oriented programming (OOP) is a way to help programmers understand program state.
-
-In our case, the `CounterV2` class is nothing more than a namespace bringing all the functions that use `CounterState` into one location. Exercise for the reader: do you think it makes sense to keep it as a class?
-
-Incidentally, you've already seen an example of this strategy in the JAX pseudo-randomness API, {mod}`jax.random`, shown in the :ref:`pseudorandom-numbers` section.
-Unlike Numpy, which manages random state using implicitly updated stateful classes, JAX requires the programmer to work directly with the random generator state -- the PRNG key.
-
-
-## Simple worked example: Linear Regression
-
-Let's apply this strategy to a simple machine learning model: linear regression via gradient descent.
-
-Here, we only deal with one kind of state: the model parameters.
-But generally, you'll see many kinds of state being threaded in and out of JAX functions, like optimizer state, layer statistics for batchnorm, and others.
-
-The function to look at carefully is `update`.
-
-```{code-cell}
-from typing import NamedTuple
-
-class Params(NamedTuple):
-  weight: jnp.ndarray
-  bias: jnp.ndarray
-
-
-def init(rng) -> Params:
-  """Returns the initial model params."""
-  weights_key, bias_key = jax.random.split(rng)
-  weight = jax.random.normal(weights_key, ())
-  bias = jax.random.normal(bias_key, ())
-  return Params(weight, bias)
-
-
-def loss(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
-  """Computes the least squares error of the model's predictions on x against y."""
-  pred = params.weight * x + params.bias
-  return jnp.mean((pred - y) ** 2)
-
-
-LEARNING_RATE = 0.005
-
-@jax.jit
-def update(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> Params:
-  """Performs one SGD update step on params using the given data."""
-  grad = jax.grad(loss)(params, x, y)
-
-  # If we were using Adam or another stateful optimizer,
-  # we would also do something like
-  #
-  #   updates, new_optimizer_state = optimizer(grad, optimizer_state)
-  # 
-  # and then use `updates` instead of `grad` to actually update the params.
-  # (And we'd include `new_optimizer_state` in the output, naturally.)
-
-  new_params = jax.tree_map(
-      lambda param, g: param - g * LEARNING_RATE, params, grad)
-
-  return new_params
-```
-
-Notice that we manually pipe the params in and out of the update function.
-
-```{code-cell}
-import matplotlib.pyplot as plt
-
-rng = jax.random.key(42)
-
-# Generate true data from y = w*x + b + noise
-true_w, true_b = 2, -1
-x_rng, noise_rng = jax.random.split(rng)
-xs = jax.random.normal(x_rng, (128, 1))
-noise = jax.random.normal(noise_rng, (128, 1)) * 0.5
-ys = xs * true_w + true_b + noise
-
-# Fit regression
-params = init(rng)
-for _ in range(1000):
-  params = update(params, xs, ys)
-
-plt.scatter(xs, ys)
-plt.plot(xs, params.weight * xs + params.bias, c='red', label='Model Prediction')
-plt.legend();
-```
-
-## Taking it further
-
-The strategy described above is how any JAX program must handle state when using transformations like `jit`, `vmap`, `grad`, etc.
-
-Handling parameters manually seems fine if you're dealing with two parameters, but what if it's a neural net with dozens of layers? You might already be getting worried about two things:
-
-1) Are we supposed to initialize them all manually, essentially repeating what we already write in the forward pass definition?
-
-2) Are we supposed to pipe all these things around manually?
-
-The details can be tricky to handle, but there are examples of libraries that take care of this for you. See [JAX Neural Network Libraries](https://github.com/google/jax#neural-network-libraries) for some examples.
diff --git a/docs/tutorials/thinking-in-jax.md b/docs/tutorials/thinking-in-jax.md
new file mode 100644
index 000000000..06cae1140
--- /dev/null
+++ b/docs/tutorials/thinking-in-jax.md
@@ -0,0 +1,458 @@
+---
+jupytext:
+  formats: md:myst
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.16.1
+kernelspec:
+  display_name: Python 3
+  language: python
+  name: python3
+---
+
+```{code-cell}
+:tags: [remove-cell]
+
+# This ensures that code cell tracebacks appearing below will be concise.
+%xmode minimal
+```
+
+(thinking-in-jax)=
+# How to think in JAX
+
+In this tutorial you will learn about how JAX operates, so that you can use it more effectively. JAX provides a simple and powerful API for writing accelerated numerical code, and working effectively in JAX sometimes requires extra consideration. This document will help you build a ground-up understanding of the JAX API.
+
+
+## JAX versus NumPy
+
+**Key concepts:**
+
+- JAX provides a NumPy-inspired interface for convenience.
+- Through [duck typing](https://en.wikipedia.org/wiki/Duck_typing), JAX arrays (`jax.Array`s) can often be used as drop-in replacements of NumPy arrays ([`ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html)s).
+- Unlike NumPy arrays, JAX arrays are always immutable.
+
+NumPy provides a well-known, powerful API for working with numerical data. For convenience, JAX provides `jax.numpy` which closely mirrors the numpy API and provides easy entry into JAX. Almost anything that can be done with `numpy` can be done with `jax.numpy`:
+
+```{code-cell}
+import matplotlib.pyplot as plt
+import numpy as np # Ordinary NumPy
+
+x_np = np.linspace(0, 10, 1000)
+y_np = 2 * np.sin(x_np) * np.cos(x_np)
+plt.plot(x_np, y_np);
+```
+
+```{code-cell}
+import jax.numpy as jnp # JAX NumPy
+
+x_jnp = jnp.linspace(0, 10, 1000)
+y_jnp = 2 * jnp.sin(x_jnp) * jnp.cos(x_jnp)
+plt.plot(x_jnp, y_jnp);
+```
+
+The code blocks are identical aside from replacing NumPy (`np`) with JAX NumPy (`jnp`), and the results are the same. JAX arrays can often be used directly in place of NumPy arrays for things like plotting.
+
+The arrays themselves are implemented as different Python types:
+
+```{code-cell}
+type(x_np)
+```
+
+```{code-cell}
+type(x_jnp)
+```
+
+Python's [duck typing](https://en.wikipedia.org/wiki/Duck_typing) allows JAX arrays `jax.Array`s and NumPy arrays `numpy.ndarray`s to be used interchangeably in many places.
+
+However, there is one important difference between JAX and NumPy arrays: JAX arrays are immutable, meaning that once created their contents cannot be changed.
+
+Here is an example of mutating an array in NumPy:
+
+```{code-cell}
+# NumPy: mutable arrays
+x = np.arange(10)
+x[0] = 10
+print(x)
+```
+
+The equivalent in JAX results in an error, as JAX arrays are immutable:
+
+```{code-cell}
+:tags: [raises-exception]
+
+# JAX: immutable arrays
+x = jnp.arange(10)
+x[0] = 10
+```
+
+For updating individual elements, JAX provides an [indexed update syntax](https://jax.readthedocs.io/en/latest/jax.ops.html#indexed-update-operators) that returns an updated copy:
+
+```{code-cell}
+y = x.at[0].set(10)
+print(x)
+print(y)
+```
+
+(thinking-in-jax-jax-arrays)=
+## JAX arrays (`jax.Array`s)
+
+**Key concepts:**
+
+- `jax.Array` is the default array implementation in JAX.
+- JAX arrays may be stored on a single device, or sharded across many devices.
+
+When you create an array in JAX, the type is `jax.Array`:
+
+```{code-cell}
+import jax
+import jax.numpy as jnp
+
+x = jnp.arange(10)
+isinstance(x, jax.Array)
+```
+
+`jax.Array` is also the appropriate type annotation for functions with array inputs or outputs:
+
+```{code-cell}
+def f(x: jax.Array) -> jax.Array:
+  return jnp.sin(x) ** 2 + jnp.cos(x) ** 2
+```
+
+JAX Array objects have a `devices` method that lets you inspect where the contents of the array are stored. In the simplest cases, this will be a single CPU device:
+
+```{code-cell}
+x.devices()
+```
+
+In general, an array may be *sharded* across multiple devices, in a manner that can be inspected via the `sharding` attribute:
+
+```{code-cell}
+x.sharding
+```
+
+In this case the sharding is on a single device, but in general a JAX array can be
+sharded across multiple devices, or even multiple hosts.
+To read more about sharded arrays and parallel computation, refer to {ref}`single-host-sharding`
+
+(thinking-in-jax-pytrees)=
+## Pytrees
+
+**Key concepts:**
+
+- JAX supports tuples, dicts, lists, and more general containers of arrays through the
+  *pytree* abstraction.
+
+Often it is convenient for applications to work with collections of arrays: for example,
+a neural network might organize its parameters in a dictionary of arrays with meaningful
+keys. Rather than handle such structures on a case-by-case basis, JAX relies on a *pytree*
+abstraction to treat such collections in a uniform matter.
+In JAX any pytree is safe to pass to transformed functions, which makes them much more flexible
+than if they only accepted single arrays as arguments.
+
+Here are some examples of objects that can be treated as pytrees:
+
+```{code-cell}
+# (nested) list of parameters
+params = [1, 2, (jnp.arange(3), jnp.ones(2))]
+
+print(jax.tree.structure(params))
+print(jax.tree.leaves(params))
+```
+
+```{code-cell}
+# Dictionary of parameters
+params = {'n': 5, 'W': jnp.ones((2, 2)), 'b': jnp.zeros(2)}
+
+print(jax.tree.structure(params))
+print(jax.tree.leaves(params))
+```
+
+```{code-cell}
+# Named tuple of parameters
+from typing import NamedTuple
+
+class Params(NamedTuple):
+  a: int
+  b: float
+
+params = Params(1, 5.0)
+print(jax.tree.structure(params))
+print(jax.tree.leaves(params))
+```
+
+JAX has a number of general-purpose utilities for working with PyTrees; for example
+the functions {func}`jax.tree.map` can be used to map a function to every leaf in a
+tree, and {func}`jax.tree.reduce` can be used to apply a reduction across the leaves
+in a tree.
+
+You can learn more in the {ref}`working-with-pytrees` tutorial.
+
+(thinking-in-jax-jax-api-layering)=
+## NumPy, `jax.lax` and XLA: JAX API layering
+
+**Key concepts:**
+
+- {mod}`jax.numpy` is a high-level wrapper that provides a familiar interface.
+- {mod}`jax.lax` is a lower-level API that is stricter and often more powerful.
+- All JAX operations are implemented in terms of operations in [XLA](https://www.tensorflow.org/xla/) — the Accelerated Linear Algebra compiler.
+
+If you look at the source of {mod}`jax.numpy`, you'll see that all the operations are eventually expressed in terms of functions defined in {mod}`jax.lax`. You can think of {mod}`jax.lax` as a stricter, but often more powerful, API for working with multi-dimensional arrays.
+
+For example, while {mod}`jax.numpy` will implicitly promote arguments to allow operations between mixed data types, {mod}`jax.lax` will not:
+
+```{code-cell}
+import jax.numpy as jnp
+jnp.add(1, 1.0)  # `jax.numpy` API implicitly promotes mixed types.
+```
+
+```{code-cell}
+:tags: [raises-exception]
+
+from jax import lax
+lax.add(1, 1.0)  # `jax.lax` API requires explicit type promotion.
+```
+
+If using {mod}`jax.lax` directly, you'll have to do type promotion explicitly in such cases:
+
+```{code-cell}
+lax.add(jnp.float32(1), 1.0)
+```
+
+Along with this strictness, {mod}`jax.lax` also provides efficient APIs for some more general operations than are supported by NumPy.
+
+For example, consider a 1D convolution, which can be expressed in NumPy this way:
+
+```{code-cell}
+x = jnp.array([1, 2, 1])
+y = jnp.ones(10)
+jnp.convolve(x, y)
+```
+
+Under the hood, this NumPy operation is translated to a much more general convolution implemented by {func}`jax.lax.conv_general_dilated`:
+
+```{code-cell}
+from jax import lax
+result = lax.conv_general_dilated(
+    x.reshape(1, 1, 3).astype(float),  # note: explicit promotion
+    y.reshape(1, 1, 10),
+    window_strides=(1,),
+    padding=[(len(y) - 1, len(y) - 1)])  # equivalent of `padding='full'`` in NumPy
+result[0, 0]
+```
+
+This is a batched convolution operation designed to be efficient for the types of convolutions often used in deep neural nets. It requires much more boilerplate, but is far more flexible and scalable than the convolution provided by NumPy (Refer to [Convolutions in JAX](https://jax.readthedocs.io/en/latest/notebooks/convolutions.html) for more details on JAX convolutions).
+
+At their heart, all {mod}`jax.lax` operations are Python wrappers for operations in XLA. Here, for example, the convolution implementation is provided by [XLA:ConvWithGeneralPadding](https://www.tensorflow.org/xla/operation_semantics#convwithgeneralpadding_convolution).
+Every JAX operation is eventually expressed in terms of these fundamental XLA operations, which is what enables just-in-time (JIT) compilation.
+
+
+(thinking-in-jax-to-jit-or-not-to-jit)=
+## To JIT or not to JIT (`jax.jit`)
+
+**Key concepts:**
+
+- By default JAX executes operations one at a time, in sequence.
+- Using a just-in-time (JIT) compilation decorator — {func}`jax.jit` — sequences of operations can be optimized together and run at once.
+- Not all JAX code can be JIT compiled, as it requires array shapes to be static & known at compile time.
+
+The fact that all JAX operations are expressed in terms of XLA allows JAX to use the XLA compiler to execute blocks of code very efficiently.
+
+For example, consider this function that normalizes the rows of a 2D matrix, expressed in terms of {mod}`jax.numpy` operations:
+
+```{code-cell}
+import jax.numpy as jnp
+
+def norm(X):
+  X = X - X.mean(0)
+  return X / X.std(0)
+```
+
+A JIT-compiled version of the function can be created using the {func}`jax.jit` transform:
+
+```{code-cell}
+from jax import jit
+norm_compiled = jit(norm)
+```
+
+This function returns the same results as the original, up to standard floating-point accuracy:
+
+```{code-cell}
+np.random.seed(1701)
+X = jnp.array(np.random.rand(10000, 10))
+np.allclose(norm(X), norm_compiled(X), atol=1E-6)
+```
+
+But due to the compilation (which includes fusing of operations, avoidance of allocating temporary arrays, and a host of other tricks), execution times can be orders of magnitude faster in the JIT-compiled case (note the use of {func}`jax.block_until_ready` to account for JAX's [asynchronous dispatch](https://jax.readthedocs.io/en/latest/async_dispatch.html)):
+
+```{code-cell}
+%timeit norm(X).block_until_ready()
+%timeit norm_compiled(X).block_until_ready()
+```
+
+That said, {func}`jax.jit` does have limitations: in particular, it requires all arrays to have static shapes. This means some JAX operations are incompatible with JIT compilation.
+
+For example, this operation can be executed in op-by-op mode:
+
+```{code-cell}
+def get_negatives(x):
+  return x[x < 0]
+
+x = jnp.array(np.random.randn(10))
+get_negatives(x)
+```
+
+But it returns an error if you attempt to execute it in {func}`jax.jit` mode:
+
+```{code-cell}
+:tags: [raises-exception]
+
+jit(get_negatives)(x)
+```
+
+This is because the function generates an array whose shape is not known at compile time: the size of the output depends on the values of the input array, and so it is not compatible with JIT.
+
+(thinking-in-jax-jit-mechanics)=
+## JIT mechanics: tracing and static variables
+
+**Key concepts:**
+
+- JIT and other JAX transforms work by *tracing* a function to determine its effect on inputs of a specific shape and type.
+
+- Variables that you don't want to be traced can be marked as *static*
+
+To use {func}`jax.jit` effectively, it is useful to understand how it works. Let's put a few `print()` statements within a JIT-compiled function and then call the function:
+
+```{code-cell}
+@jit
+def f(x, y):
+  print("Running f():")
+  print(f"  x = {x}")
+  print(f"  y = {y}")
+  result = jnp.dot(x + 1, y + 1)
+  print(f"  result = {result}")
+  return result
+
+x = np.random.randn(3, 4)
+y = np.random.randn(4)
+f(x, y)
+```
+
+Notice that the print statements execute, but rather than printing the data you passed to the function, though, it prints *tracer* objects that stand-in for them.
+
+These tracer objects are what {func}`jax.jit` uses to extract the sequence of operations specified by the function. Basic tracers are stand-ins that encode the **shape** and **dtype** of the arrays, but are agnostic to the values. This recorded sequence of computations can then be efficiently applied within XLA to new inputs with the same shape and dtype, without having to re-execute the Python code.
+
+When you call the compiled function again on matching inputs, no recompilation is required and nothing is printed because the result is computed in compiled XLA rather than in Python:
+
+```{code-cell}
+x2 = np.random.randn(3, 4)
+y2 = np.random.randn(4)
+f(x2, y2)
+```
+
+The extracted sequence of operations is encoded in a JAX expression, or *jaxpr* for short. You can view the jaxpr using the {func}`jax.make_jaxpr` transformation:
+
+```{code-cell}
+from jax import make_jaxpr
+
+def f(x, y):
+  return jnp.dot(x + 1, y + 1)
+
+make_jaxpr(f)(x, y)
+```
+
+Note one consequence of this: because JIT compilation is done *without* information on the content of the array, control flow statements in the function cannot depend on traced values. For example, this fails:
+
+```{code-cell}
+:tags: [raises-exception]
+
+@jit
+def f(x, neg):
+  return -x if neg else x
+
+f(1, True)
+```
+
+If there are variables that you would not like to be traced, they can be marked as static for the purposes of JIT compilation:
+
+```{code-cell}
+from functools import partial
+
+@partial(jit, static_argnums=(1,))
+def f(x, neg):
+  return -x if neg else x
+
+f(1, True)
+```
+
+Note that calling a JIT-compiled function with a different static argument results in recompilation, so the function still works as expected:
+
+```{code-cell}
+f(1, False)
+```
+
+Understanding which values and operations will be static and which will be traced is a key part of using {func}`jax.jit` effectively.
+
+
+(thinking-in-jax-static-versus-traced-operations)=
+## Static versus traced operations
+
+**Key concepts:**
+
+- Just as values can be either static or traced, operations can be static or traced.
+- Static operations are evaluated at compile-time in Python. Traced operations are compiled & evaluated at run-time in XLA.
+- Use NumPy (`numpy`) for operations that you want to be static. Use JAX NumPy {mod}`jax.numpy` for operations that you want to be traced.
+
+This distinction between static and traced values makes it important to think about how to keep a static value static. Consider this function:
+
+```{code-cell}
+:tags: [raises-exception]
+
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def f(x):
+  return x.reshape(jnp.array(x.shape).prod())
+
+x = jnp.ones((2, 3))
+f(x)
+```
+
+This fails with an error specifying that a tracer was found instead of a 1D sequence of concrete values of integer type. Let's add some print statements to the function to understand why this is happening:
+
+```{code-cell}
+@jit
+def f(x):
+  print(f"x = {x}")
+  print(f"x.shape = {x.shape}")
+  print(f"jnp.array(x.shape).prod() = {jnp.array(x.shape).prod()}")
+  # comment this out to avoid the error:
+  # return x.reshape(jnp.array(x.shape).prod())
+
+f(x)
+```
+
+Notice that although `x` is traced, `x.shape` is a static value. However, when you use {func}`jnp.array` and {func}`jnp.prod` on this static value, it becomes a traced value, at which point it cannot be used in a function like `reshape()` that requires a static input (recall: array shapes must be static).
+
+A useful pattern is to:
+
+- Use NumPy (`numpy`) for operations that should be static (i.e., done at compile-time); and
+- Use JAX NumPy (`jax.numpy`) for operations that should be traced (i.e. compiled and executed at run-time).
+
+For this function, it might look like this:
+
+```{code-cell}
+from jax import jit
+import jax.numpy as jnp
+import numpy as np
+
+@jit
+def f(x):
+  return x.reshape((np.prod(x.shape),))
+
+f(x)
+```
+
+For this reason, a standard convention in JAX programs is to `import numpy as np` and `import jax.numpy as jnp` so that both interfaces are available for finer control over whether operations are performed in a static matter (with `numpy`, once at compile-time) or a traced manner (with {mod}`jax.numpy`, optimized at run-time).
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 92a27d5fb..03df5f231 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -376,8 +376,7 @@ def define_bool_state(
   an error.
   """
   if not isinstance(default, bool):
-    raise TypeError(f"Default value must be of type bool, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type bool, got {default}")
   default = bool_env(name.upper(), default)
   name = name.lower()
   if upgrade:
@@ -421,8 +420,7 @@ def define_enum_state(
     A contextmanager to control the thread-local state value.
   """
   if not isinstance(default, str):
-    raise TypeError(f"Default value must be of type str, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type str, got {default}")
   name = name.lower()
   default = os.getenv(name.upper(), default)
   if default not in enum_values:
@@ -478,8 +476,7 @@ def define_optional_enum_state(
     A contextmanager to control the thread-local state value.
   """
   if default is not None and not isinstance(default, str):
-    raise TypeError(f"Default value must be of type str or None, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type str or None, got {default}")
   name = name.lower()
   default = os.getenv(name.upper(), default)
   if default is not None and default not in enum_values:
@@ -529,8 +526,7 @@ def define_int_state(
     A contextmanager to control the thread-local state value.
   """
   if not isinstance(default, int):
-    raise TypeError(f"Default value must be of type int, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type int, got {default}")
   name = name.lower()
   default_env = os.getenv(name.upper())
   if default_env is not None:
@@ -576,8 +572,7 @@ def define_float_state(
     A contextmanager to control the thread-local state value.
   """
   if not isinstance(default, float):
-    raise TypeError(f"Default value must be of type float, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type float, got {default}")
   name = name.lower()
   default_env = os.getenv(name.upper())
   if default_env is not None:
@@ -629,8 +624,7 @@ def define_string_state(
     A contextmanager to control the thread-local state value.
   """
   if not isinstance(default, str):
-    raise TypeError(f"Default value must be of type str, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type str, got {default}")
 
   def validator(new_val):
     if not isinstance(new_val, str):
@@ -673,8 +667,7 @@ def define_optional_string_state(
     A contextmanager to control the thread-local state value.
   """
   if default is not None and not isinstance(default, str):
-    raise TypeError(f"Default value must be of type str or None, got {default} "
-                    f"of type {getattr(type(default), '__name__', type(default))}")
+    raise TypeError(f"Default value must be of type str or None, got {default}")
 
   def validator(new_val):
     if new_val is not None and not isinstance(new_val, str):
diff --git a/jax/_src/debugger/web_debugger.py b/jax/_src/debugger/web_debugger.py
index 443bfa676..d5c97d190 100644
--- a/jax/_src/debugger/web_debugger.py
+++ b/jax/_src/debugger/web_debugger.py
@@ -13,9 +13,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import atexit
-import functools
-import importlib.util
 import os
 from typing import Any
 import weakref
@@ -23,22 +20,16 @@ import weakref
 from jax._src.debugger import cli_debugger
 from jax._src.debugger import core as debugger_core
 
-
-@functools.cache
-def _web_pdb_version() -> tuple[int, ...]:
+web_pdb_version: tuple[int, ...] | None = None
+try:
   import web_pdb  # pytype: disable=import-error
-  return tuple(map(int, web_pdb.__version__.split(".")))
-
-
-_web_consoles: dict[tuple[str, int], Any] = {}
+  web_pdb_version = tuple(map(int, web_pdb.__version__.split(".")))
+  WEB_PDB_ENABLED = True
+except:
+  WEB_PDB_ENABLED = False
 
 
-@atexit.register
-def _close_debuggers():
-  for console in _web_consoles.values():
-    console.close()
-  _web_consoles.clear()
-
+_web_consoles: dict[tuple[str, int], web_pdb.WebConsole] = {}
 
 class WebDebugger(cli_debugger.CliDebugger):
   """A web-based debugger."""
@@ -48,7 +39,6 @@ class WebDebugger(cli_debugger.CliDebugger):
   def __init__(self, frames: list[debugger_core.DebuggerFrame], thread_id,
                completekey: str = "tab", host: str = "", port: int = 5555):
     if (host, port) not in _web_consoles:
-      import web_pdb  # pytype: disable=import-error
       _web_consoles[host, port] = web_pdb.WebConsole(host, port, self)
     # Clobber the debugger in the web console
     _web_console = _web_consoles[host, port]
@@ -64,7 +54,7 @@ class WebDebugger(cli_debugger.CliDebugger):
     current_line = None
     if current_frame.offset is not None:
       current_line = current_frame.offset + 1
-    if _web_pdb_version() < (1, 4, 4):
+    if web_pdb_version and web_pdb_version < (1, 4, 4):
       return {
         'filename': filename,
         'listing': '\n'.join(lines),
@@ -84,15 +74,15 @@ class WebDebugger(cli_debugger.CliDebugger):
 
   def get_globals(self):
     current_frame = self.current_frame()
-    return "\n".join(
-        f"{key} = {value}"
-        for key, value in sorted(current_frame.globals.items()))
+    globals = "\n".join([f"{key} = {value}" for key, value in
+      sorted(current_frame.globals.items())])
+    return globals
 
   def get_locals(self):
     current_frame = self.current_frame()
-    return "\n".join(
-        f"{key} = {value}"
-        for key, value in sorted(current_frame.locals.items()))
+    locals = "\n".join([f"{key} = {value}" for key, value in
+      sorted(current_frame.locals.items())])
+    return locals
 
   def run(self):
     return self.cmdloop()
@@ -101,6 +91,5 @@ def run_debugger(frames: list[debugger_core.DebuggerFrame],
                  thread_id: int | None, **kwargs: Any):
   WebDebugger(frames, thread_id, **kwargs).run()
 
-
-if importlib.util.find_spec("web_pdb") is not None:
+if WEB_PDB_ENABLED:
   debugger_core.register_debugger("web", run_debugger, -2)
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index e923d1e5f..84fc677c9 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import importlib.util
 from collections.abc import Sequence
 import functools
 import string
@@ -45,6 +44,20 @@ from jax._src.lib.mlir.dialects import hlo
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
 
+# pytype: disable=import-error
+try:
+  import rich
+  import rich.align
+  import rich.box
+  import rich.console
+  import rich.padding
+  import rich.style
+  import rich.table
+  RICH_ENABLED = True
+except:
+  RICH_ENABLED = False
+# pytype: enable=import-error
+
 class DebugEffect(effects.Effect):
   __str__ = lambda self: "Debug"
 debug_effect = DebugEffect()
@@ -427,19 +440,8 @@ def visualize_sharding(shape: Sequence[int], sharding: Sharding, *,
                        min_width: int = 9, max_width: int = 80,
                        color_map: ColorMap | None = None):
   """Visualizes a ``Sharding`` using ``rich``."""
-  if not importlib.util.find_spec("rich"):
+  if not RICH_ENABLED:
     raise ValueError("`visualize_sharding` requires `rich` to be installed.")
-
-  # These imports are local so that they don't affect JAX import times.
-  # pytype: disable=import-error
-  import rich.align
-  import rich.console
-  import rich.box
-  import rich.padding
-  import rich.style
-  import rich.table
-  # pytype: enable=import-error
-
   if len(shape) > 2 or len(shape) < 1:
     raise ValueError(
         "`visualize_sharding` only works for shapes with 1 and 2 dimensions.")
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index 85dd03ee3..2667176d2 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2037,12 +2037,6 @@ def to_gspmd_sharding(s: sharding_impls.XLACompatibleSharding,
                          memory_kind=s.memory_kind)
 
 
-# Dummy function which is a no-op in OSS since enhanced barrier is switched on
-# in OSS.
-def spmd_mode_check(da_object, inline):
-  return
-
-
 @profiler.annotate_function
 def lower_sharding_computation(
     closed_jaxpr: core.ClosedJaxpr,
@@ -2133,7 +2127,19 @@ def lower_sharding_computation(
       da_object,
       it.chain(in_shardings, out_shardings, [js for js, _ in jaxpr_sharding]))  # type: ignore
 
-  spmd_mode_check(da_object, inline)
+  if not da_object.is_fully_addressable:  # type: ignore
+    if inline and config.spmd_mode.value != 'allow_all':
+      raise RuntimeError(
+          "Running operations on `Array`s that are not fully addressable by this "
+          "process (i.e. `Array`s with data sharded across multiple devices and "
+          "processes.) is dangerous. It’s very important that all processes run "
+          "the same cross-process computations in the same order otherwise it "
+          "can lead to hangs. "
+          "If you’re not already familiar with JAX’s multi-process "
+          "programming model, please read "
+          "https://jax.readthedocs.io/en/latest/multi_process.html. "
+          "To fix this error, run your `jitted` computation inside "
+          "`with jax.spmd_mode('allow_all'):` context manager.")
 
   # 2. Build up the HLO
   semantic_in_shardings = SemanticallyEqualShardings(
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 8c5594cb4..0e628fc5e 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -26,7 +26,7 @@ import numpy as np
 
 from jax import lax
 from jax._src import api
-from jax._src import core, config
+from jax._src import core
 from jax._src import dtypes
 from jax._src.numpy import ufuncs
 from jax._src.numpy.util import (
@@ -708,42 +708,6 @@ nancumsum = _make_cumulative_reduction(np.nancumsum, lax.cumsum,
 nancumprod = _make_cumulative_reduction(np.nancumprod, lax.cumprod,
                                         fill_nan=True, fill_value=1)
 
-@implements(getattr(np, 'cumulative_sum', None))
-def cumulative_sum(
-    x: ArrayLike, /, *, axis: int | None = None,
-    dtype: DTypeLike | None = None,
-    include_initial: bool = False) -> Array:
-  check_arraylike("cumulative_sum", x)
-  x = lax_internal.asarray(x)
-  if x.ndim == 0:
-    raise ValueError(
-      "The input must be non-scalar to take a cumulative sum, however a "
-      "scalar value or scalar array was given."
-    )
-  if axis is None:
-    axis = 0
-    if x.ndim > 1:
-      raise ValueError(
-        f"The input array has rank {x.ndim}, however axis was not set to an "
-        "explicit value. The axis argument is only optional for one-dimensional "
-        "arrays.")
-
-  axis = _canonicalize_axis(axis, x.ndim)
-  dtypes.check_user_dtype_supported(dtype)
-  kind = x.dtype.kind
-  if (dtype is None and kind in {'i', 'u'}
-      and x.dtype.itemsize*8 < int(config.default_dtype_bits.value)):
-    dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
-  x = x.astype(dtype=dtype or x.dtype)
-  out = cumsum(x, axis=axis)
-  if include_initial:
-    zeros_shape = list(x.shape)
-    zeros_shape[axis] = 1
-    out = lax_internal.concatenate(
-      [lax_internal.full(zeros_shape, 0, dtype=out.dtype), out],
-      dimension=axis)
-  return out
-
 # Quantiles
 @implements(np.quantile, skip_params=['out', 'overwrite_input'])
 @partial(api.jit, static_argnames=('axis', 'overwrite_input', 'interpolation',
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index f06cd99ba..27ac106a4 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -1718,8 +1718,8 @@ lowering_rules[for_loop.for_p] = _for_lowering_rule
 
 
 def _lower_jaxpr_to_for_loop(ctx: LoweringRuleContext,
-                             jaxpr: jax_core.Jaxpr, start: int | ir.Value,
-                             num_steps: int | ir.Value, consts, *args,
+                             jaxpr: jax_core.Jaxpr, start: int,
+                             num_steps: int, consts, *args,
                              has_loop_index: bool,
                              unroll: int):
   def _run_body(i, args):
@@ -1736,11 +1736,7 @@ def _lower_jaxpr_to_for_loop(ctx: LoweringRuleContext,
       args = jaxpr_subcomp(lowering_context, jaxpr, *consts, *args)
     return args
 
-  if (
-      not isinstance(start, ir.Value)
-      and not isinstance(num_steps, ir.Value)
-      and num_steps == unroll
-  ):
+  if num_steps == unroll:
     # No need for an scf.For. We can just unroll completely
     for i in range(start, start + num_steps):
       args = _run_body(
@@ -1751,9 +1747,13 @@ def _lower_jaxpr_to_for_loop(ctx: LoweringRuleContext,
   if unroll != 1:
     raise NotImplementedError(
         f"Only unroll={num_steps=} and unroll=1 supported. Got {unroll=}.")
-  i32 = jax_core.ShapedArray((), jnp.int32)
-  lbd = _ensure_mlir_value(start, i32)
-  ubd = arith.addi(lbd, _ensure_mlir_value(num_steps, i32))
+  lbd = ir_constant(0, mlir_type=mlir.dtype_to_ir_type(jnp.dtype("int32")))
+  if isinstance(num_steps, int):
+    ubd = ir_constant(
+        num_steps, mlir_type=_dtype_to_ir_type(jnp.dtype("int32"))
+    )
+  else:
+    ubd = num_steps
   step = ir_constant(1, mlir_type=_dtype_to_ir_type(jnp.dtype("int32")))
   for_op = scf.ForOp(lbd, ubd, step, args)
   with ir.InsertionPoint(for_op.body):
@@ -1792,7 +1792,7 @@ def _scan_lowering_rule(
     linear: tuple[bool, ...],
     length: int,
     reverse: bool,
-    unroll: bool | int,
+    unroll: bool,
     num_consts: int,
     num_carry: int,
     _split_transpose: bool,
@@ -1808,9 +1808,9 @@ def _scan_lowering_rule(
   if jaxpr_consts: raise NotImplementedError
   del jaxpr_consts
 
-  jaxpr, has_loop_index = pallas_utils.pattern_match_scan_to_fori_loop(
-      jaxpr, num_consts, num_carry
-  )
+  jaxpr, has_loop_index = (
+      pallas_utils.pattern_match_scan_to_fori_loop(jaxpr, num_consts, num_carry)
+      )
   consts, args = split_list(args, [num_consts])
   consts_avals, args_avals = split_list(ctx.avals_in, [num_consts])
   if has_loop_index:
@@ -1856,7 +1856,7 @@ def _while_lowering_rule(
       ),
       jaxpr,
       lb,
-      arith.subi(ub, lb),
+      ub,
       body_consts,
       *args,
       has_loop_index=True,
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index a921de565..9d7c40e4a 100644
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -451,7 +451,7 @@ def ndtri(p: ArrayLike) -> Array:
   to `p`.
 
   A piece-wise rational approximation is done for the function.
-  This is based on the implementation in netlib.
+  This is a based on the implementation in netlib.
 
   Args:
     p: an array of type `float32`, `float64`.
@@ -829,7 +829,7 @@ def bessel_jn(z: ArrayLike, *, v: int, n_iter: int=50) -> Array:
 def _gen_recurrence_mask(
     l_max: int, is_normalized: bool, dtype: Any
 ) -> tuple[Array, Array]:
-  """Generates a mask for recurrence relation on the remaining entries.
+  """Generates mask for recurrence relation on the remaining entries.
 
   The remaining entries are with respect to the diagonal and offdiagonal
   entries.
@@ -984,7 +984,7 @@ def _gen_associated_legendre(l_max: int,
   `Y_l^m(θ, φ) = N_l^m * P_l^m(cos(θ)) * exp(i m φ)`, where `N_l^m` is the
   normalization factor and θ and φ are the colatitude and longitude,
   respectively. `N_l^m` is chosen in the way that the spherical harmonics form
-  a set of orthonormal basis functions of L^2(S^2). For the computational
+  a set of orthonormal basis function of L^2(S^2). For the computational
   efficiency of spherical harmonics transform, the normalization factor is
   used in the computation of the ALFs. In addition, normalizing `P_l^m`
   avoids overflow/underflow and achieves better numerical stability. Three
@@ -1008,7 +1008,7 @@ def _gen_associated_legendre(l_max: int,
       operation, `W` is a diagonal matrix containing the quadrature weights,
       and `I` is the identity matrix. The Gauss-Chebyshev points are equally
       spaced, which only provide approximate discrete orthogonality. The
-      Driscoll & Healy quadrature points are equally spaced and provide the
+      Driscoll & Healy qudarture points are equally spaced and provide the
       exact discrete orthogonality. The number of sampling points is required to
       be twice as the number of frequency points (modes) in the Driscoll & Healy
       approach, which enables FFT and achieves a fast spherical harmonics
@@ -1219,7 +1219,7 @@ def sph_harm(m: Array,
 
   Args:
     m: The order of the harmonic; must have `|m| <= n`. Return values for
-      `|m| > n` are undefined.
+      `|m| > n` ara undefined.
     n: The degree of the harmonic; must have `n >= 0`. The standard notation for
       degree in descriptions of spherical harmonics is `l (lower case L)`. We
       use `n` here to be consistent with `scipy.special.sph_harm`. Return
@@ -1229,7 +1229,7 @@ def sph_harm(m: Array,
     n_max: The maximum degree `max(n)`. If the supplied `n_max` is not the true
       maximum value of `n`, the results are clipped to `n_max`. For example,
       `sph_harm(m=jnp.array([2]), n=jnp.array([10]), theta, phi, n_max=6)`
-      actually returns
+      acutually returns
       `sph_harm(m=jnp.array([2]), n=jnp.array([6]), theta, phi, n_max=6)`
   Returns:
     A 1D array containing the spherical harmonics at (m, n, theta, phi).
@@ -1709,7 +1709,7 @@ def spence(x: Array) -> Array:
     -\int_0^z \frac{\log(1 - t)}{t}dt
     \end{equation}
 
-  This is our spence(1 - z).
+  this is our spence(1 - z).
   """
   x = jnp.asarray(x)
   dtype = lax.dtype(x)
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
index 0c7ab97f9..339c9ea00 100644
--- a/jax/_src/xla_bridge.py
+++ b/jax/_src/xla_bridge.py
@@ -106,11 +106,9 @@ _CPU_ENABLE_GLOO_COLLECTIVES = config.DEFINE_bool(
     help="If True, enable cross-process collectives on CPU using Gloo.",
 )
 
-# TODO(yueshengys): turn default back to True after resolving memory increase
-# issue.
 _CPU_ENABLE_ASYNC_DISPATCH = config.DEFINE_bool(
     name="jax_cpu_enable_async_dispatch",
-    default=False,
+    default=True,
     help="Only applies to non-parallel computations. If False, run computations"
     "inline without async dispatch.",
 )
diff --git a/jax/experimental/array_api/__init__.py b/jax/experimental/array_api/__init__.py
index 7050a4e3c..dfb2ff988 100644
--- a/jax/experimental/array_api/__init__.py
+++ b/jax/experimental/array_api/__init__.py
@@ -204,7 +204,6 @@ from jax.experimental.array_api._sorting_functions import (
 )
 
 from jax.experimental.array_api._statistical_functions import (
-    cumulative_sum as cumulative_sum,
     max as max,
     mean as mean,
     min as min,
diff --git a/jax/experimental/array_api/_statistical_functions.py b/jax/experimental/array_api/_statistical_functions.py
index 141b80abf..2e1333317 100644
--- a/jax/experimental/array_api/_statistical_functions.py
+++ b/jax/experimental/array_api/_statistical_functions.py
@@ -18,10 +18,6 @@ from jax.experimental.array_api._data_type_functions import (
 )
 
 
-def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False):
-  """Calculates the cumulative sum of elements in the input array x."""
-  return jax.numpy.cumulative_sum(x, axis=axis, dtype=dtype, include_initial=include_initial)
-
 def max(x, /, *, axis=None, keepdims=False):
   """Calculates the maximum value of the input array x."""
   return jax.numpy.max(x, axis=axis, keepdims=keepdims)
diff --git a/jax/experimental/array_api/skips.txt b/jax/experimental/array_api/skips.txt
index c502f4ac5..76149272a 100644
--- a/jax/experimental/array_api/skips.txt
+++ b/jax/experimental/array_api/skips.txt
@@ -1,13 +1,11 @@
 # Known failures for the array api tests.
 
 # Test suite attempts in-place mutation:
+array_api_tests/test_special_cases.py::test_binary
 array_api_tests/test_special_cases.py::test_iop
 array_api_tests/test_special_cases.py::test_nan_propagation
-array_api_tests/test_array_object.py::test_setitem
-
-# Raises NonInteractiveExampleWarning
-array_api_tests/test_special_cases.py::test_binary
 array_api_tests/test_special_cases.py::test_unary
+array_api_tests/test_array_object.py::test_setitem
 
 # fft test suite is buggy as of 83f0bcdc
 array_api_tests/test_fft.py
diff --git a/jax/extend/core/__init__.py b/jax/extend/core/__init__.py
index dbe53f7c2..2732b1984 100644
--- a/jax/extend/core/__init__.py
+++ b/jax/extend/core/__init__.py
@@ -19,4 +19,15 @@ from jax._src.abstract_arrays import (
   array_types as array_types
 )
 
+from jax._src.core import (
+  ClosedJaxpr as ClosedJaxpr,
+  Jaxpr as Jaxpr,
+  JaxprEqn as JaxprEqn,
+  jaxpr_as_fun as jaxpr_as_fun,
+  Literal as Literal,
+  Primitive as Primitive,
+  Token as Token,
+  Var as Var,
+)
+
 from . import primitives as primitives
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
index 1b9a990f3..0b8a8cdd0 100644
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -296,7 +296,6 @@ from jax._src.numpy.reductions import (
     count_nonzero as count_nonzero,
     cumsum as cumsum,
     cumprod as cumprod,
-    cumulative_sum as cumulative_sum,
     max as max,
     mean as mean,
     median as median,
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
index 274063804..7c7e68a06 100644
--- a/jax/numpy/__init__.pyi
+++ b/jax/numpy/__init__.pyi
@@ -241,9 +241,6 @@ def cumprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
 cumproduct = cumprod
 def cumsum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...) -> Array: ...
-def cumulative_sum(x: ArrayLike, /, *, axis: int | None = ...,
-                   dtype: DTypeLike | None = ...,
-                   include_initial: bool = ...) -> Array: ...
 
 def deg2rad(x: ArrayLike, /) -> Array: ...
 degrees = rad2deg
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
index a321c3b8f..749eedb14 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -1000,185 +1000,6 @@ LogicalResult scf_for_rule(RewriteContext &ctx, Operation &op,
   return success();
 }
 
-LogicalResult scf_while_rule(RewriteContext &ctx, Operation &op,
-                             const ArrayRef<Layout> layouts_in,
-                             const ArrayRef<Layout> layouts_out) {
-  scf::WhileOp while_op = cast<scf::WhileOp>(op);
-  TPU_ASSERT_EQ_OP(layouts_in.size(), while_op->getNumOperands());
-  TPU_ASSERT_EQ_OP(layouts_out.size(), while_op->getNumResults());
-  TPU_ASSERT_EQ_OP(layouts_in.size(), layouts_out.size());
-
-  // The terminator for the before region is the condition op.
-  // It takes multiple arguments -- the first being the decision to execute the
-  // after region or branch to the exit.
-  FAILUREOR_ASSIGN_OR_RETURN(
-      const SmallVector<Layout> condition_in_layouts,
-      getInLayouts(*while_op.getBeforeBody()->getTerminator(),
-                   ctx.target_shape));
-  if (!llvm::equal(ArrayRef<Layout>(condition_in_layouts).drop_front(1),
-                   layouts_out)) {
-    return op.emitOpError(
-        "Mismatched layouts between scf.while result and its before region "
-        "condition.");
-  }
-
-  if (failed(applyLayoutBlock(ctx, *while_op.getBeforeBody()))) {
-    return failure();
-  }
-
-  FAILUREOR_ASSIGN_OR_RETURN(
-      const SmallVector<Layout> after_yield_in_layouts,
-      getInLayouts(*while_op.getYieldOp(), ctx.target_shape));
-  if (!layouts_out.empty() &&
-      ArrayRef<Layout>(after_yield_in_layouts) != layouts_out) {
-    return op.emitOpError(
-        "Not implemented: different layouts while's yield's operands and  "
-        "results");
-  }
-
-  if (failed(applyLayoutBlock(ctx, *while_op.getAfterBody()))) {
-    return failure();
-  }
-
-  if (op.getNumResults() == 0) {
-    return success();
-  }
-
-  OpBuilder builder(&op);
-  SmallVector<Value> unrolled_args;
-  for (int i = 0; i < layouts_in.size(); ++i) {
-    auto layout = layouts_in[i];
-    auto operand = while_op.getOperand(i);
-    if (auto vector_operand = dyn_cast<TypedValue<VectorType>>(operand)) {
-      if (!layout.has_value()) {
-        return op.emitOpError("Expected layout for vector operand");
-      }
-      FAILUREOR_ASSIGN_OR_RETURN(
-          const xla::Array<Value> tiles,
-          disassemble(builder, *layout, vector_operand, ctx.target_shape));
-      unrolled_args.append(tiles.begin(), tiles.end());
-    } else {
-      if (layout.has_value()) {
-        return op.emitOpError("Expected no layout for scalar operand");
-      }
-      unrolled_args.push_back(operand);
-    }
-  }
-
-  // Create a new scf::WhileOp with unrolled args.
-  auto new_op = builder.create<scf::WhileOp>(
-      while_op->getLoc(),
-      TypeRange(while_op.getConditionOp().getOperands().drop_front(1)),
-      unrolled_args, nullptr, nullptr);
-
-  const auto tile_body_args = [&](::mlir::Block *old_body,
-                                  ::mlir::Block *new_body,
-                                  const ArrayRef<Layout> layouts) {
-    TPU_ASSERT_OP(old_body != nullptr);
-    TPU_ASSERT_OP(new_body != nullptr);
-    int num_old_args = old_body->getNumArguments();
-    SmallVector<Location> locs(new_body->getNumArguments(), while_op.getLoc());
-    old_body->addArguments(TypeRange(new_body->getArguments()), locs);
-    builder.setInsertionPointToStart(old_body);
-    auto arg_idx = num_old_args;
-    for (auto [old_arg, layout] : llvm::zip_equal(
-             old_body->getArguments().take_front(num_old_args), layouts)) {
-      if (const auto vty = dyn_cast<VectorType>(old_arg.getType())) {
-        TPU_ASSERT_OP(layout.has_value());
-        const SmallVector<int64_t> tiles_shape =
-            layout->tileArrayShape(vty.getShape(), ctx.target_shape);
-        const int64_t num_vectors = ShapedType::getNumElements(tiles_shape);
-        xla::Array<Value> tiles(tiles_shape);
-        TPU_ASSERT_LE_OP(arg_idx + num_vectors, old_body->getNumArguments());
-        tiles.SetValues(llvm::make_range(
-            old_body->getArguments().begin() + arg_idx,
-            old_body->getArguments().begin() + arg_idx + num_vectors));
-        arg_idx += num_vectors;
-        RollVectorsOp rolled_op =
-            assemble(builder, vty, *layout, tiles, ctx.target_shape);
-        old_arg.replaceUsesWithIf(rolled_op, [&](OpOperand &operand) {
-          return operand.getOwner() != rolled_op;
-        });
-      } else {
-        TPU_ASSERT_OP(!layout.has_value());
-        old_arg.replaceAllUsesWith(old_body->getArgument(arg_idx));
-        ++arg_idx;
-      }
-    }
-    old_body->eraseArguments(0, num_old_args);
-    return success();
-  };
-
-  const auto before_status = tile_body_args(while_op.getBeforeBody(),
-                                            new_op.getBeforeBody(), layouts_in);
-  if (before_status.failed()) return before_status;
-  new_op.getBefore().takeBody(while_op.getBefore());
-
-  const auto after_status = tile_body_args(while_op.getAfterBody(),
-                                           new_op.getAfterBody(), layouts_out);
-  if (after_status.failed()) return after_status;
-  new_op.getAfter().takeBody(while_op.getAfter());
-
-  builder.setInsertionPointAfter(new_op);
-  int64_t res_idx = 0;
-  SmallVector<Value> rolled_results;
-  for (auto [result, layout] :
-       llvm::zip_equal(while_op.getResults(), layouts_out)) {
-    if (const auto vty = dyn_cast<VectorType>(result.getType())) {
-      TPU_ASSERT_OP(layout.has_value());
-      const SmallVector<int64_t> tiles_shape =
-          layout->tileArrayShape(vty.getShape(), ctx.target_shape);
-      const int64_t num_vectors = ShapedType::getNumElements(tiles_shape);
-      xla::Array<Value> tiles(tiles_shape);
-      TPU_ASSERT_LE_OP(res_idx + num_vectors, new_op.getResults().size());
-      tiles.SetValues(llvm::make_range(
-          new_op.getResults().begin() + res_idx,
-          new_op.getResults().begin() + res_idx + num_vectors));
-      res_idx += num_vectors;
-      RollVectorsOp rolled_op =
-          assemble(builder, vty, *layout, tiles, ctx.target_shape);
-      rolled_results.push_back(rolled_op);
-    } else {
-      TPU_ASSERT_OP(!layout.has_value());
-      rolled_results.push_back(new_op.getResult(res_idx));
-      ++res_idx;
-    }
-  }
-
-  while_op.replaceAllUsesWith(rolled_results);
-  while_op.erase();
-  return success();
-}
-
-LogicalResult scf_condition_rule(RewriteContext &ctx, Operation &op,
-                                 const ArrayRef<Layout> layouts_in,
-                                 const ArrayRef<Layout> layouts_out) {
-  OpBuilder builder(&op);
-  auto condition_op = cast<scf::ConditionOp>(op);
-  TPU_ASSERT_EQ_OP(layouts_in.size(), condition_op.getNumOperands());
-  TPU_ASSERT_EQ_OP(layouts_out.size(), 0);
-  SmallVector<Value> unrolled;
-
-  for (auto [operand, layout] :
-       llvm::zip_equal(condition_op.getOperands(), layouts_in)) {
-    if (auto vector_operand = dyn_cast<TypedValue<VectorType>>(operand)) {
-      // When the operand has vector type, disassemble the operand.
-      TPU_ASSERT_OP(layout.has_value());
-      FAILUREOR_ASSIGN_OR_RETURN(
-          const xla::Array<Value> tiles,
-          disassemble(builder, *layout, vector_operand, ctx.target_shape));
-      unrolled.append(tiles.begin(), tiles.end());
-    } else {
-      TPU_ASSERT_OP(!layout.has_value());
-      unrolled.push_back(operand);
-    }
-  }
-
-  // Replace the old operands with unrolled operands.
-  condition_op->setOperands(unrolled);
-  return success();
-}
-
 LogicalResult scf_if_rule(RewriteContext &ctx, Operation &op,
                           const ArrayRef<Layout> layouts_in,
                           const ArrayRef<Layout> layouts_out) {
@@ -3813,8 +3634,6 @@ const llvm::StringMap<rule_type> &rules() {
       {arith::TruncIOp::getOperationName(), arith_trunci_rule},
       {func::ReturnOp::getOperationName(), func_return_rule},
       {scf::ForOp::getOperationName(), scf_for_rule},
-      {scf::WhileOp::getOperationName(), scf_while_rule},
-      {scf::ConditionOp::getOperationName(), scf_condition_rule},
       {scf::IfOp::getOperationName(), scf_if_rule},
       {scf::YieldOp::getOperationName(), scf_yield_rule},
       {tpu::RotateOp::getOperationName(), tpu_rotate_rule},
@@ -4367,33 +4186,63 @@ FailureOr<TypedValue<VectorType>> relayout(
         });
     src = dst;
     src_tiles = std::move(src_tiles_retiled);
-  } else if (  // TODO(b/265133506): Generalize retiling.
-               // (8,128) -> (8 * packing,128) tiling change for packed type.
+  } else if (  // TODO(b/265133506): Generalize retiling to general 16-bit types
+               // (might need to use a different unpacking op).
+               // (8,128) -> (16,128) tiling change for packed 16-bit types.
+      src.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+      dst.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+      vty.getElementTypeBitWidth() == 16 && src.offsets() == dst.offsets() &&
+      src.tiling() == std::array<int64_t, 2>{8, 128} &&
+      dst.tiling() == std::array<int64_t, 2>{16, 128}) {
+    const VectorLayout new_src(src.bitwidth(), src.offsets(), dst.tiling());
+    xla::Array<Value> src_tiles_retiled(
+        new_src.tileArrayShape(vty.getShape(), target_shape));
+    src_tiles_retiled.Each([&](absl::Span<const int64_t> idx, Value *tile) {
+      SmallVector<int64_t> src_idx(idx.begin(), idx.end());
+      src_idx[src_idx.size() - 2] *= 2;
+      src_idx[src_idx.size() - 1] /= 2;
+      Value src_row1 = src_tiles(src_idx);
+      if (src_idx[src_idx.size() - 2] + 1 <
+          src_tiles.dim(src_tiles.num_dimensions() - 2)) {
+        ++src_idx[src_idx.size() - 2];
+      }
+      Value src_row2 = src_tiles(src_idx);
+      const int vreg_part = idx[idx.size() - 1] % 2;
+
+      VectorType vreg_x32 =
+          vty.getElementType().isSignlessInteger()
+              ? VectorType::get(target_shape, builder.getI32Type())
+              : VectorType::get(target_shape, builder.getF32Type());
+      auto half_row1 = builder.create<tpu::UnpackSubelementsOp>(
+          v.getLoc(), vreg_x32, src_row1, vreg_part);
+      auto half_row2 = builder.create<tpu::UnpackSubelementsOp>(
+          v.getLoc(), vreg_x32, src_row2, vreg_part);
+      *tile = builder.create<tpu::PackSubelementsOp>(
+          v.getLoc(), src_row1.getType(), ValueRange{half_row1, half_row2});
+    });
+    src = new_src;
+    src_tiles = std::move(src_tiles_retiled);
+  } else if (  // (8,128) -> (32,128) tiling change for packed 8-bit integers.
       src.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
       dst.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
-      vty.getElementTypeBitWidth() < 32 &&
-      32 % vty.getElementTypeBitWidth() == 0 &&
+      vty.getElementType() == builder.getI8Type() &&
       src.offsets() == dst.offsets() &&
       src.tiling() == std::array<int64_t, 2>{8, 128} &&
-      dst.tiling() == std::array<int64_t, 2>{8 * dst.packing(), 128}) {
+      dst.tiling() == std::array<int64_t, 2>{32, 128}) {
     const VectorLayout new_src(src.bitwidth(), src.offsets(), dst.tiling());
     xla::Array<Value> src_tiles_retiled(
         new_src.tileArrayShape(vty.getShape(), target_shape));
-    int vty_packing = dst.packing();
-    VectorType vreg_x32 =
-        vty.getElementType().isSignlessInteger()
-            ? VectorType::get(target_shape, builder.getI32Type())
-            : VectorType::get(target_shape, builder.getF32Type());
+    VectorType vreg_i32 =
+        getNativeVregType(builder.getI32Type(), target_shape).value();
     src_tiles_retiled.Each([&](absl::Span<const int64_t> idx, Value *tile) {
-      const int vreg_part = idx.back() % vty_packing;
-      SmallVector<Value, 8> parts;
-      parts.reserve(vty_packing);
+      const int vreg_part = idx.back() % 4;
+      std::array<Value, 4> parts;
       SmallVector<int64_t> src_idx(idx.begin(), idx.end());
-      src_idx[src_idx.size() - 2] *= vty_packing;
-      src_idx[src_idx.size() - 1] /= vty_packing;
-      for (int i = 0; i < vty_packing; ++i) {
-        parts.push_back(builder.create<tpu::UnpackSubelementsOp>(
-            v.getLoc(), vreg_x32, src_tiles(src_idx), vreg_part));
+      src_idx[src_idx.size() - 2] *= 4;
+      src_idx[src_idx.size() - 1] /= 4;
+      for (int i = 0; i < 4; ++i) {
+        parts[i] = builder.create<tpu::UnpackSubelementsOp>(
+            v.getLoc(), vreg_i32, src_tiles(src_idx), vreg_part);
         if (src_idx[src_idx.size() - 2] <
             src_tiles.dim(src_tiles.num_dimensions() - 2) - 1) {
           ++src_idx[src_idx.size() - 2];
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
index 1469ba2cc..a69bcc3d8 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -220,14 +220,6 @@ class VectorLayoutInferer {
         if (infer(op).failed()) {
           return failure();
         }
-      } else if (auto op = dyn_cast<scf::WhileOp>(any_op)) {
-        if (infer(op).failed()) {
-          return failure();
-        }
-      } else if (auto op = dyn_cast<scf::ConditionOp>(any_op)) {
-        if (infer(op).failed()) {
-          return failure();
-        }
       } else if (auto op = dyn_cast<tpu::RotateOp>(any_op)) {
         if (infer(op).failed()) {
           return failure();
@@ -544,118 +536,6 @@ class VectorLayoutInferer {
     return success();
   }
 
-  LogicalResult infer(scf::WhileOp op) {
-    static LogicalResult (*match_condition)(Operation *) = [](Operation *op) {
-      TPU_CHECK_OP(isa<scf::ConditionOp>(op), "expected condition terminator");
-      return success();
-    };
-    static LogicalResult (*match_yield)(Operation *) = [](Operation *op) {
-      TPU_CHECK_OP(isa<scf::YieldOp>(op), "expected yield terminator");
-      return success();
-    };
-    TPU_CHECK_OP(op.getNumRegions() == 2, "expected two blocks for scf.while");
-
-    const auto layout_for_type = [&op, this](const ::mlir::Value &arg,
-                                             SmallVector<Layout> *layouts) {
-      if (arg.getType().isSignlessIntOrIndexOrFloat()) {
-        layouts->push_back(kNoLayout);
-      } else if (isa<VectorType>(arg.getType())) {
-        auto layout = getLayout(arg);
-        layouts->push_back(layout);
-      } else {
-        op.emitOpError() << "unsupported arg type " << arg.getType()
-                         << " in scf.while";
-        return failure();
-      }
-      return success();
-    };
-
-    SmallVector<Layout> in_layouts;
-    in_layouts.reserve(op->getNumOperands());
-    for (const auto &arg : op.getInits()) {
-      const auto status = layout_for_type(arg, &in_layouts);
-      if (status.failed()) return status;
-    }
-
-    // Formally, the types and layouts of the results should follow the layout
-    // of the condition op in the Before region, rather than mimicking the input
-    // layouts. In practice these are constrained to be the same for our current
-    // pipelines, but doesn't represent the full expressiveness of scf.while.
-    // TODO(hmckenzie): Base output layout on ConditionOp, not inputs.
-    SmallVector<Layout> out_layouts = in_layouts;
-
-    // Use tpu.assume_layout to annotate every block argument with the layout of
-    // the corresponding operand in WhileOp and replace all uses of the block
-    // argument with the result of tpu.assume_layout.
-    ImplicitLocOpBuilder builder =
-        ImplicitLocOpBuilder::atBlockBegin(op.getLoc(), op.getBeforeBody());
-    for (auto [iter_arg, layout] :
-         llvm::zip_equal(op.getBeforeBody()->getArguments(), in_layouts)) {
-      if (!dyn_cast<VectorType>(iter_arg.getType())) {
-        continue;
-      }
-      auto assume_layout_op =
-          builder.create<AssumeLayoutOp>(iter_arg.getType(), iter_arg);
-      setLayout(assume_layout_op, layout, layout);
-      iter_arg.replaceUsesWithIf(assume_layout_op, [&](OpOperand &operand) {
-        return operand.getOwner() != assume_layout_op;
-      });
-    }
-    if (inferBlock(*op.getBeforeBody(), match_condition).failed()) {
-      return failure();
-    }
-
-    builder =
-        ImplicitLocOpBuilder::atBlockBegin(op.getLoc(), op.getAfterBody());
-    for (auto [iter_arg, layout] :
-         llvm::zip_equal(op.getAfterBody()->getArguments(), out_layouts)) {
-      if (!dyn_cast<VectorType>(iter_arg.getType())) {
-        continue;
-      }
-      auto assume_layout_op =
-          builder.create<AssumeLayoutOp>(iter_arg.getType(), iter_arg);
-      setLayout(assume_layout_op, layout, layout);
-      iter_arg.replaceUsesWithIf(assume_layout_op, [&](OpOperand &operand) {
-        return operand.getOwner() != assume_layout_op;
-      });
-    }
-
-    if (inferBlock(*op.getAfterBody(), match_yield).failed()) {
-      return failure();
-    }
-
-    auto *condition_op = op.getBeforeBody()->getTerminator();
-    SmallVector<Layout> cond_layout;
-    cond_layout.reserve(out_layouts.size() + 1);
-    cond_layout.push_back(kNoLayout);
-    cond_layout.append(out_layouts);
-    setInLayout(condition_op, cond_layout);
-
-    auto *yield_op = op.getAfterBody()->getTerminator();
-    setInLayout(yield_op, in_layouts);
-
-    setLayout(op, in_layouts, out_layouts);
-    return success();
-  }
-  LogicalResult infer(scf::ConditionOp op) {
-    SmallVector<Layout> in_layouts;
-    in_layouts.reserve(op->getNumOperands());
-    for (const auto &arg : op.getOperands()) {
-      if (arg.getType().isSignlessIntOrIndexOrFloat()) {
-        in_layouts.push_back(kNoLayout);
-      } else if (isa<VectorType>(arg.getType())) {
-        auto layout = getLayout(arg);
-        in_layouts.push_back(layout);
-      } else {
-        op.emitOpError() << "unsupported arg type " << arg.getType()
-                         << " in scf::condition";
-        return failure();
-      }
-    }
-    setLayout(op, in_layouts, ArrayRef<Layout>(in_layouts).drop_front(1));
-    return success();
-  }
-
   LogicalResult infer(tpu::RotateOp op) {
     auto bitwidth = op.getType().getElementTypeBitWidth();
     if (bitwidth != 32) {
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index 9e09f07d6..9b9ba0332 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -20,8 +20,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "ca98bc910ffa778e25ca784c6201a431dfb383d4"
-XLA_SHA256 = "f78196224303453d8e68539bfbb3ec2e713e748da162de72a1f36dde947296c4"
+XLA_COMMIT = "6805a043c6b834fa9b3edc4ff953f10d1fa1f386"
+XLA_SHA256 = "798f66030e38f96accca04c8fe8c6fb1e1e04b060b578128bf7f7044faa6295f"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index 5667c3459..91c64e749 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -68,7 +68,6 @@ MAIN_NAMESPACE = {
   'copysign',
   'cos',
   'cosh',
-  'cumulative_sum',
   'divide',
   'e',
   'empty',
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index 73100352c..b72e47d9d 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -770,64 +770,5 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     self.assertAllClose(expected, actual, atol=0)
 
 
-  @jtu.sample_product(
-    [dict(shape=shape, axis=axis)
-      for shape in all_shapes
-      for axis in list(
-        range(-len(shape), len(shape))
-      ) + ([None] if len(shape) == 1 else [])],
-    dtype=all_dtypes + [None],
-    out_dtype=all_dtypes,
-    include_initial=[False, True],
-  )
-  @jtu.ignore_warning(category=NumpyComplexWarning)
-  @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises mixed type promotion
-  def testCumulativeSum(self, shape, axis, dtype, out_dtype, include_initial):
-    rng = jtu.rand_some_zero(self.rng())
-
-    def np_mock_op(x, axis=None, dtype=None, include_initial=False):
-      kind = x.dtype.kind
-      if (dtype is None and kind in {'i', 'u'}
-          and x.dtype.itemsize*8 < int(config.default_dtype_bits.value)):
-        dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
-      axis = axis or 0
-      x = x.astype(dtype=dtype or x.dtype)
-      out = jnp.cumsum(x, axis=axis)
-      if include_initial:
-        zeros_shape = list(x.shape)
-        zeros_shape[axis] = 1
-        out = jnp.concat([jnp.zeros(zeros_shape, dtype=out.dtype), out], axis=axis)
-      return out
-
-
-    # We currently "cheat" to ensure we have JAX arrays, not NumPy arrays as
-    # input because we rely on JAX-specific casting behavior
-    args_maker = lambda: [jnp.array(rng(shape, dtype))]
-    np_op = getattr(np, "cumulative_sum", np_mock_op)
-    kwargs = dict(axis=axis, dtype=out_dtype, include_initial=include_initial)
-
-    np_fun = lambda x: np_op(x, **kwargs)
-    jnp_fun = lambda x: jnp.cumulative_sum(x, **kwargs)
-    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
-
-
-  @jtu.sample_product(
-      shape=filter(lambda x: len(x) != 1, all_shapes), dtype=all_dtypes,
-      include_initial=[False, True])
-  def testCumulativeSumErrors(self, shape, dtype, include_initial):
-    rng = jtu.rand_some_zero(self.rng())
-    x = rng(shape, dtype)
-    rank = jnp.asarray(x).ndim
-    if rank == 0:
-      msg = r"The input must be non-scalar to take"
-      with self.assertRaisesRegex(ValueError, msg):
-        jnp.cumulative_sum(x, include_initial=include_initial)
-    elif rank > 1:
-      msg = r"The input array has rank \d*, however"
-      with self.assertRaisesRegex(ValueError, msg):
-        jnp.cumulative_sum(x, include_initial=include_initial)
-
-
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/pallas/pallas_call_tpu_test.py b/tests/pallas/pallas_call_tpu_test.py
index cf352149d..8268ff775 100644
--- a/tests/pallas/pallas_call_tpu_test.py
+++ b/tests/pallas/pallas_call_tpu_test.py
@@ -1699,17 +1699,16 @@ class PallasCallWhileLoopTest(PallasTPUTest):
         return i < j
 
       def body(carry):
-        io, j = carry
-        i = io - 128
+        i, j = carry
         sl = sl = jax.lax.div(i, 128)
         l = jax.lax.rem(i, 128)
         v = x_ref[0, sl, l]
         s = pl.load(r_ref, (0, 0))
         pl.store(r_ref, (0, 0), s + v)
-        return io + 1, j
+        return i + 1, j
 
-      i = 128
-      j = 128 + 1024
+      i = 0
+      j = 1024
       i, j = jax.lax.while_loop(cond, body, (i, j))
 
     x = jnp.arange(4096)
@@ -1731,26 +1730,6 @@ class PallasCallWhileLoopTest(PallasTPUTest):
     expected = jnp.sum(jnp.arange(1024))
     np.testing.assert_array_equal(r, expected)
 
-  def test_fori(self):
-    """Tests lowering of a while_loop which can reduce to a fori_loop."""
-
-    def kernel(lb_ref, ub_ref, o_ref):
-      o_ref[0, 0] = 0
-
-      def body(i, _):
-        o_ref[0, 0] += 1
-
-      jax.lax.fori_loop(lb_ref[0, 0], ub_ref[0, 0], body, None)
-
-    smem = pl.BlockSpec(memory_space=pltpu.SMEM)
-    r = pl.pallas_call(
-        kernel,
-        in_specs=(smem, smem),
-        out_specs=smem,
-        out_shape=jax.ShapeDtypeStruct([1, 1], jnp.int32),
-    )(*(jnp.array([[x]]) for x in (2, 6)))
-    np.testing.assert_array_equal(r, 4)
-
 
 class PallasCallPipelineTest(parameterized.TestCase):
 
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index 10326c3dd..7e8f8f56d 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -1367,8 +1367,9 @@ class IOCallbackTest(jtu.JaxTestCase):
         return i + 1
       jax.lax.while_loop(lambda i: i < 2, body, 0)
 
-    jax.vmap(f)(jnp.arange(3.))  # don't crash
+    jax.vmap(f)(jnp.arange(3.))
     jax.effects_barrier()
+    self.assertAllClose(x_lst, [0., 1., 2., 0., 2., 4.] * 2, check_dtypes=False)
 
 
 if __name__ == "__main__":