jax Documentation

File: README.md Copy Chapter

Transformable numerical computing at scale

Transformations
| Scaling
| Install guide
| Change logs
| Reference docs

What is JAX? Copy Chapter

JAX is a Python library for accelerator-oriented array computation and program transformation,
designed for high-performance numerical computing and large-scale machine learning.

JAX can automatically differentiate native
Python and NumPy functions. It can differentiate through loops, branches,
recursion, and closures, and it can take derivatives of derivatives of
derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation)
via jax.grad as well as forward-mode differentiation,
and the two can be composed arbitrarily to any order.

JAX uses XLA
to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators.
You can compile your own pure functions with jax.jit.
Compilation and automatic differentiation can be composed arbitrarily.

Dig a little deeper, and you'll see that JAX is really an extensible system for
composable function transformations at scale.

This is a research project, not an official Google product. Expect
sharp edges.
Please help by trying it out, reporting bugs,
and letting us know what you think!

import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.tanh(outputs)  # inputs to the next layer
  return outputs                # no activation on last layer

def loss(params, inputs, targets):
  preds = predict(params, inputs)
  return jnp.sum((preds - targets)**2)

grad_loss = jax.jit(jax.grad(loss))  # compiled gradient evaluation function
perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))  # fast per-example grads

Contents Copy Section

• Transformations
• Scaling
• Current gotchas
• Installation
• Citing JAX
• Reference documentation

Transformations Copy Chapter

At its core, JAX is an extensible system for transforming numerical functions.
Here are three: jax.grad, jax.jit, and jax.vmap.

Automatic differentiation with `grad` Copy Section

Use jax.grad
to efficiently compute reverse-mode gradients:

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743

You can differentiate to any order with grad:

print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
# prints 0.62162673

You're free to use differentiation with Python control flow:

def abs_val(x):
  if x > 0:
    return x
  else:
    return -x

abs_val_grad = jax.grad(abs_val)
print(abs_val_grad(1.0))   # prints 1.0
print(abs_val_grad(-1.0))  # prints -1.0 (abs_val is re-evaluated)

See the JAX Autodiff
Cookbook
and the reference docs on automatic
differentiation
for more.

Compilation with `jit` Copy Section

Use XLA to compile your functions end-to-end with
jit,
used either as an @jit decorator or as a higher-order function.

import jax
import jax.numpy as jnp

def slow_f(x):
  # Element-wise ops see a large benefit from fusion
  return x * x + x * 2.0

x = jnp.ones((5000, 5000))
fast_f = jax.jit(slow_f)
%timeit -n10 -r3 fast_f(x)
%timeit -n10 -r3 slow_f(x)

Using jax.jit constrains the kind of Python control flow
the function can use; see
the tutorial on Control Flow and Logical Operators with JIT
for more.

Auto-vectorization with `vmap` Copy Section

vmap maps
a function along array axes.
But instead of just looping over function applications, it pushes the loop down
onto the function’s primitive operations, e.g. turning matrix-vector multiplies into
matrix-matrix multiplies for better performance.

Using vmap can save you from having to carry around batch dimensions in your
code:

import jax
import jax.numpy as jnp

def l1_distance(x, y):
  assert x.ndim == y.ndim == 1  # only works on 1D inputs
  return jnp.sum(jnp.abs(x - y))

def pairwise_distances(dist1D, xs):
  return jax.vmap(jax.vmap(dist1D, (0, None)), (None, 0))(xs, xs)

xs = jax.random.normal(jax.random.key(0), (100, 3))
dists = pairwise_distances(l1_distance, xs)
dists.shape  # (100, 100)

By composing jax.vmap with jax.grad and jax.jit, we can get efficient
Jacobian matrices, or per-example gradients:

per_example_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)))

Scaling Copy Chapter

To scale your computations across thousands of devices, you can use any
composition of these:

• Compiler-based automatic parallelization
  where you program as if using a single global machine, and the compiler chooses
  how to shard data and partition computation (with some user-provided constraints);
• Explicit sharding and automatic partitioning
  where you still have a global view but data shardings are
  explicit in JAX types, inspectable using jax.typeof;
• Manual per-device programming
  where you have a per-device view of data
  and computation, and can communicate with explicit collectives.

from jax.sharding import set_mesh, AxisType, PartitionSpec as P
mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

# parameters are sharded for FSDP:
for W, b in params:
  print(f'{jax.typeof(W)}')  # f32[512@data,512]
  print(f'{jax.typeof(b)}')  # f32[512]

# shard data for batch parallelism:
inputs, targets = jax.device_put((inputs, targets), P('data'))

# evaluate gradients, automatically parallelized!
gradfun = jax.jit(jax.grad(loss))
param_grads = gradfun(params, (inputs, targets))

See the tutorial and
advanced guides for more.

Gotchas and sharp bits Copy Chapter

See the Gotchas
Notebook.

Installation Copy Chapter

Supported platforms Copy Section

Instructions Copy Section

See the documentation
for information on alternative installation strategies. These include compiling
from source, installing with Docker, using other versions of CUDA, a
community-supported conda build, and answers to some frequently-asked questions.

Citing JAX Copy Chapter

To cite this repository:

@software{jax2018github,
  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
  url = {http://github.com/jax-ml/jax},
  version = {0.3.13},
  year = {2018},
}

In the above bibtex entry, names are in alphabetical order, the version number
is intended to be that from jax/version.py, and
the year corresponds to the project's open-source release.

A nascent version of JAX, supporting only automatic differentiation and
compilation to XLA, was described in a paper that appeared at SysML
2018. We're currently working on
covering JAX's ideas and capabilities in a more comprehensive and up-to-date
paper.

Reference documentation Copy Chapter

For details about the JAX API, see the
reference documentation.

For getting started as a JAX developer, see the
developer documentation.

File: ci/envs/README.md Copy Chapter

JAXCI Environment Variables

This docpage describes the various JAXCI environment variables that are used
in the CI scripts and their behaviors. These variables are used to control the
behavior of the CI scripts such as the Python version used, path to JAX/XLA
repo, if to clone XLA repo, etc.

Docker Specific Environment Variables Copy Chapter

[!NOTE]
The following environment variables only affect the build if the
run_docker_container.sh
script was invoked to start a Docker container and the build is running inside
that container. Typically, this would be the internal CI builds and local
builds. Note that while GitHub actions use the same Docker images, they do not
invoke "run_docker_container.sh" as they leverage built-in containerization
features to run jobs within a container.

File: ci/utilities/README.md Copy Chapter

JAX CI Utility Scripts

This docpage gives a brief overview of the different utility scripts and what
they are used for.

• setup_build_environment.sh: Sets up the build environment such as
  cloning the latest XLA, adjusting file paths (for Windows), etc.
• convert_msys_paths_to_win_paths.py: Converts MSYS Linux-like paths
  stored in env variables to Windows paths.
• install_wheels_locally.sh: Used by Pytest scripts to install JAX wheels
  and any additional extras on the system.
• run_auditwheel.sh: Verifies that the Linux artifacts are "manylinux"
  compliant.
• run_docker_container.sh: Runs a Docker container called "jax". Images
  are read from the JAXCI_DOCKER_IMAGE environment variable in
  ci/envs/docker.env.

File: ci/README.md Copy Chapter

JAX Continuous Integration

This folder contains the configuration files and scripts used to build and test
JAX. It is typically used by continuous integration (CI) jobs to automate builds
and run comprehensive tests across various platforms and configurations. This
page provides an overview of the JAX CI system, its components, and the
different workflows it supports.

JAX's CI System Copy Chapter

JAX's CI system is composed of several interacting components and orchestrates
builds and tests using a hybrid approach, leveraging both an internal CI system
and GitHub Actions as well as an internal build orchestrator for managing
nightly and release flows. It encompasses several distinct workflows, including
comprehensive presubmit checks triggered on pull requests and branch pushes,
bi-hourly continuous builds, extensive nightly builds with broad platform
coverage, and a controlled release process that culminates in PyPI publication.

These flows build four packages: jax, jaxlib, jax-cuda-plugin,
jax-cuda-pjrt and support a range of environments, including:

• Linux x86: CPU, TPU, CUDA
• Linux aarch64: CPU, CUDA
• Windows x86: CPU
• Mac Arm64: CPU

Architecture Overview Copy Section

1. Internal CI System: An internal CI system is used for specific build and
   test tasks, such as nightly builds, release candidate (RC) builds, and
   Mac-specific testing.

2. GitHub Actions: Used for presubmit checks, continuous integration builds
   and tests, and nightly/release artifact testing.

3. Build Orchestrator: An internal tool used to manage complex workflows
   such as nightly / release flows, promoting RC builds to release, etc.

4. Artifact Storage:

• Google Cloud Storage (GCS) Buckets: Used for temporary storage of artifacts
  between jobs in GitHub Actions workflows and for storing packages built
  during nightly and release flows before testing.
• Artifact Registry: Used to store nightly packages, RC packages and final
  releases.
• PyPI: Where final releases are published.

CI Workflows and Where They Run Copy Section

JAX's CI system consists of the following workflows:

1. Presubmits: Presubmits are run in GitHub actions and are triggered on
   pull requests that target the main branch and on pushes to the main and
   release branch. JAX's presubmit run time SLO is about 10 minutes so these
   are typically run using Bazel with remote build execution
   (RBE). RBE allows us to execute build and
   test actions on a distributed system, separate from the local machine,
   instead of solely on the local machine. This enables faster build and test
   times by utilizing parallel computing resources and caching across a cluster
   of machines. However, we also use Pytest in workflows where we are not able
   to use RBE such as the TPU presubmit. In such presubmits, we usually run a
   subset of tests to be able to satisfy the presubmit run time SLO. To see the
   list of the presubmit workflows,
   click here.

2. Continuous: These jobs are run in GitHub actions and are scheduled to
   run once every 2 hours on the main branch. It builds JAX packages and runs
   a wide range of tests targeting different environments such as CPU, CUDA
   (L4, H100, B200, etc), and TPU (v4-8, v5e-8, etc.). For more information,
   see
   wheel_tests_continuous.yml
   (An example run.)

3. Nightly Builds and Tests: These jobs use an hybrid approach of both the
   internal CI system and GitHub actions. The jobs are triggered once every
   night by the internal build orchestrator tool. It first triggers the jobs in
   the internal CI system to build the JAX packages for different
   configurations (Python versions, CUDA versions, etc) and uploads them to a
   staging bucket in GCS as well as to the nightly artifact registry. Next,
   testing jobs are triggered that download the artifacts from the staging
   bucket and run tests. Mac testing jobs are run in the internal CI system.
   For non-Mac testing, a trigger job is run that invokes the
   wheel_tests_nightly_release.yml
   workflow in GitHub Actions. JAX's nightly artifacts can be found here:
   jax,
   jaxlib,
   jax-cuda-plugin,
   jax-cuda-pjrt.

4. Release Builds and Tests: Release flow is similar to the nightly flow
   except for few differences. First, release process has to be triggered
   manually in the internal build orchestrator and should be done only after a
   release branch (E.g release/0.5.3) has been created. The build jobs build
   two sets of artifacts for each package: 1. RC wheels 2. Final version
   wheels. These two sets are pretty much the same package except for their
   metadata and wheel tags. The RC wheels are then uploaded to the staging
   bucket and release artifact registry. After the uploads are done, the test
   jobs are triggered. As with the nightly flow, Mac test jobs are run in the
   internal CI system while non-Mac test jobs are run in GitHub actions. To see
   the GitHub actions run for a particular release, filter the workflow runs by
   its branch name.

5. Promote RC to Final and Publish to PyPI: If the RC wheels pass all
   testing, then we are ready to promote it as the final version and publish it
   to PyPI. This entire flow is internal and is run in our internal CI system.
   Final version of the packages are published to PyPI and JAX's release
   artifact registry. JAX's release artifacts (RC and final versions) can be
   found here:
   jax,
   jaxlib,
   jax-cuda-plugin,
   jax-cuda-pjrt.

JAX's Official CI and Build/Test Scripts Copy Section

JAX's CI jobs (both internal and those on GitHub actions) run the scripts in
this folder. An overview of the different folders and their purpose is given
below:

• ci/: Contains all build scripts, environment files, and utility scripts.
• ci/utilities/: Contains helper scripts used throughout the build/test
  process. See
  README.md
  for a brief overview of these utility scripts and their behavior.
• ci/envs/: Holds environment files that set JAXCI environment variables
  that control build and test configurations. see
  README.md to
  see the complete list of these variables and their behavior.

Every build script in this folder first source the JAXCI envs in
default.env and
then run the
setup_build_environment.sh
script to set up the build environment.

A brief overview of each build script in this folder is given below:

[!NOTE]
Both internal and GitHub action jobs run under the
ml-build
Docker image which contains build tools such as Python, Bazelisk, LLVM/Clang,
manylinux compliant libraries (in Linux images), etc.

• build_artifacts.sh: These build the various JAX artifacts. We build
  three different type of artifacts based on the type of job: Nightly,
  RC/Release, or at HEAD.
• run_bazel_test_cpu_rbe.sh/run_bazel_test_cuda_rbe.sh: These run Bazel
  tests with RBE on every GitHub PR. We test compatibility with both CPU and
  CUDA. On platforms where RBE is not natively supported (e.g Linux Arm64), we
  cross-compile the test targets for Linux Aarch64 on Linux x86. As the tests
  still need to be run on the host machines and because running the tests on a
  single machine can take a long time, we skip running them on these
  platforms.
  Note for run_bazel_test_cpu_rbe.sh:
  • If $JAXCI_BUILD_JAXLIB=false and $JAXCI_BUILD_JAX=false, these jobs
    depend on local JAX wheels and therefore require that the following wheels
    to be present in the ../dist folder: jax, and jaxlib wheels. In CI
    builds, we first build these wheels from source and then run the
    bazel test command.
  • If $JAXCI_BUILD_JAXLIB=false and $JAXCI_BUILD_JAX=true, CPU jobs
    depend on local jaxlib wheels and therefore require that jaxlib wheel to
    be present in the ../dist folder. GPU obs
    depend on local jaxlib and CUDA wheels, and therefore require that the
    following wheels to be present in the ../dist folder: jaxlib,
    jax-cuda-plugin, and jax-cuda-pjrt wheels. In CI builds, we first
    build these wheels from source and then run the bazel test command.
  • If $JAXCI_BUILD_JAXLIB=wheel and $JAXCI_BUILD_JAX=wheel, the Bazel
    tests use
    py_import.
  • If $JAXCI_BUILD_JAXLIB=true and $JAXCI_BUILD_JAX=true, Bazel will use
    individual targets in the test dependencies.
• run_bazel_test_cuda_non_rbe.sh: These run the following Bazel CUDA
  tests: Single accelerator tests with one GPU apiece and Multi-accelerator
  tests with all GPUs.
  • If $JAXCI_BUILD_JAXLIB=false and $JAXCI_BUILD_JAX=false, these jobs
    depend on local JAX wheels and therefore require that the following wheels
    to be present in the ../dist folder: jax, jaxlib, jax-cuda-plugin,
    and jax-cuda-pjrt wheels. In CI builds, we first build these wheels from
    source and then run the bazel test command.
  • If $JAXCI_BUILD_JAXLIB=wheel and $JAXCI_BUILD_JAX=wheel, the Bazel
    tests use py_import.
• run_pytest_*.sh: These run tests with Pytests and use the JAX wheel
  packages installed on the system. In CI builds, we build the wheels first
  from source and then run the pytest commands. We test compatibility with
  CPU, CUDA, and TPU. These are primarily run as part of the continuous and
  nightly/release test jobs except for TPU which is also run as a presubmit
  testing a subset of the tests.

Different Test Configurations Copy Chapter

JAX's CI Test jobs run under different test configurations. These configurations
are described briefly in the sections below.

XLA Versions Copy Section

JAX's CI builds rely on XLA, but use different versions depending on the type of
build. To ensure stability and reproducibility, nightly and release builds use a
pinned XLA version defined in MODULE.bazel.

However, to keep JAX compatible with the latest XLA developments, presubmit and
postsubmit builds utilize the most recent XLA version. This is done by
overriding the default XLA dependency with a local copy of the XLA repository.
We do this by passing --override_module=xla=/path/to/local/xla which
instructs Bazel to depend on the XLA in the local system instead of the pinned
version in MODULE.bazel.

The CI system uses the JAXCI environment variables to manage this process.
When running jobs that need to use XLA at head, we set JAXCI_CLONE_MAIN_XLA=1.
This clones the XLA repository at head and sets JAXCI_XLA_GIT_DIR to its path.
JAX build CLI
automatically adds the necessary Bazel flag (--override_repository) to point
to this local XLA version during the build process if JAXCI_XLA_GIT_DIR is
set. In jobs where the build CLI is not used such as the RBE presubmits, we
explicitly include --override_repository=xla="${JAXCI_XLA_GIT_DIR}" as part
of the test command.

Enabling/Disabling 64-bit Data Types Copy Section

By default, JAX enforces single-precision numbers to mitigate the Numpy API’s
tendency to aggressively promote operands to double. In order to use
double-precision numbers, we need to set the JAX_ENABLE_X64 environment
variable. In CI, we test both configurations in presubmits and postsubmits by
using the JAXCI_ENABLE_X64 environment variable.

[Googlers Only] Connecting to CI Runners for Debugging Copy Chapter

If you are a Googler, you can connect to one of the self-hosted runners we have
on GitHub to debug your workflow. For more information, see
go/ml-github-actions:connect.

Running These Scripts Locally on Your Machine Copy Chapter

[!IMPORTANT]
If you are a Linux / Windows user, you need to have Docker installed as a
prerequisite. Additionally, if running on Windows, please run these commands
in a bash environment as all the scripts are written in Shell.

Follow the steps below to run a CI script locally on your machine.

1. [Optional] Set JAXCI variables in your shell environment. See
   ci/envs/README.md
   for the list of JAXCI variables and their behavior.

2. [Linux/Windows]

   Start the Docker container by running:

   bash Copy

   ./ci/utilities/run_docker_container.sh

   This will start a Docker container named "jax". Note that if you set any
   JAXCI variables in step 1, they will also be be set in the container.

   Run the script under the Docker container.

   bash Copy

   # docker exec jax 
       docker exec jax ./ci/build_artifacts.sh jaxlib

3. [Mac] Execute the build script directly.

   bash Copy

   # ./
       ./ci/build_artifacts.sh jaxlib

File: cloud_tpu_colabs/README.md Copy Chapter

JAX on Cloud TPU examples

The same JAX code that runs on CPU and GPU can also be run on TPU. Cloud TPUs
have the advantage of quickly giving you access to multiple TPU accelerators,
including in Colab. All of the
example notebooks here use
jax.pmap to run JAX
computation across multiple TPU cores from Colab. You can also run the same code
directly on a Cloud TPU
VM.

Example Cloud TPU notebooks Copy Chapter

The following notebooks showcase how to use and what you can do with Cloud TPUs on Colab:

[Pmap Cookbook](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Pmap_Cookbook.ipynb) Copy Section

A guide to getting started with pmap, a transform for easily distributing SPMD
computations across devices.

[Lorentz ODE Solver](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb) Copy Section

Contributed by Alex Alemi (alexalemi@)

Solve and plot parallel ODE solutions with pmap.

[Wave Equation](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Wave_Equation.ipynb) Copy Section

Contributed by Stephan Hoyer (shoyer@)

Solve the wave equation with pmap, and make cool movies! The spatial domain is partitioned across the 8 cores of a Cloud TPU.

[JAX Demo](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/JAX_demo.ipynb) Copy Section

An overview of JAX presented at the Program Transformations for ML workshop at NeurIPS 2019 and the Compilers for ML workshop at CGO 2020. Covers basic numpy usage, grad, jit, vmap, and pmap.

Performance notes Copy Chapter

The guidance on running TensorFlow on TPUs applies to JAX as well, with the exception of TensorFlow-specific details. Here we highlight a few important details that are particularly relevant to using TPUs in JAX.

Padding Copy Section

One of the most common culprits for surprisingly slow code on TPUs is inadvertent padding:

• Arrays in the Cloud TPU are tiled. This entails padding one of the dimensions to a multiple of 8, and a different dimension to a multiple of 128.
• The matrix multiplication unit performs best with pairs of large matrices that minimize the need for padding.

bfloat16 dtype Copy Section

By default*, matrix multiplication in JAX on TPUs uses bfloat16 with float32 accumulation. This can be controlled with the precision keyword argument on relevant jax.numpy functions (matmul, dot, einsum, etc). In particular:

• precision=jax.lax.Precision.DEFAULT: uses mixed bfloat16 precision (fastest)
• precision=jax.lax.Precision.HIGH: uses multiple MXU passes to achieve higher precision
• precision=jax.lax.Precision.HIGHEST: uses even more MXU passes to achieve full float32 precision

JAX also adds the bfloat16 dtype, which you can use to explicitly cast arrays to bfloat16, e.g., jax.numpy.array(x, dtype=jax.numpy.bfloat16).

* We might change the default precision in the future, since it is arguably surprising. Please comment/vote on this issue if it affects you!

Running JAX on a Cloud TPU VM Copy Chapter

Refer to the Cloud TPU VM
documentation.

Reporting issues and getting help Copy Chapter

If you run into Cloud TPU-specific issues (e.g. trouble creating a Cloud TPU
VM), please email [email protected], or if
you are a TRC member. You can also file a
JAX issue or ask a discussion
question for any issues with these
notebooks or using JAX in general.

If you have any other questions or comments regarding JAX on Cloud TPUs, please
email [email protected]. We’d like to hear from you!

File: docs/_tutorials/advanced-compilation.md Copy Chapter

Advanced compilation

This is a placeholder for a section in the new {ref}`jax-tutorials-draft`.

For the time being, you may find some related content in the old documentation:
- {doc}`../aot`
- {doc}`../pallas/index`

File: docs/_tutorials/advanced-debugging.md Copy Chapter

jupytext:
formats: md:myst
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.16.4
kernelspec:
display_name: Python 3
language: python

name: python3 Copy Chapter

(advanced-debugging)=

Advanced debugging

This is a placeholder for a section in the new {ref}`jax-tutorials-draft`.

For the time being, you may find some related content in the old documentation:
- {doc}`../debugging`

File: docs/_tutorials/index.rst Copy Chapter

:orphan:

.. _jax-tutorials-draft:

JAX tutorials draft

.. note::

The tutorials below are a work in progress; for the time being, please refer
to the older tutorial content, including :ref:beginner-guide,
:ref:jax-101, and the now-deleted JAX 101 tutorials.

JAX 101 Copy Chapter

Mostly finalized at :ref:jax-101!

.. toctree::
:maxdepth: 1

../key-concepts
../jit-compilation
../automatic-vectorization
../automatic-differentiation
../debugging
../random-numbers
../pytrees
../parallel
../stateful-computations
simple-neural-network

JAX 201 Copy Chapter

.. toctree::
:maxdepth: 1

parallelism
advanced-debugging
profiling-and-performance

JAX 301 Copy Chapter

.. toctree::
:maxdepth: 1

advanced-compilation

File: docs/_tutorials/parallelism.md Copy Chapter

Parallel computation

This is a placeholder for a section in the new {ref}`jax-tutorials-draft`.

For the time being, you may find some related content in the old documentation:
- {doc}`../multi_process`
- {doc}`../parallel`

File: docs/_tutorials/profiling-and-performance.md Copy Chapter

Profiling and performance

This is a placeholder for a section in the new {ref}`jax-tutorials-draft`.

For the time being, you may find some related content in the old documentation:
- {doc}`../profiling`
- {doc}`../device_memory_profiling`
- {doc}`../transfer_guard`