JAX

https://docs.jax.dev/en/latest/
Admin
JAX is a high-performance numerical computing library that provides array manipulation, automatic...
Tokens:3,769,655
Snippets:8,816
Trust Score:6.4
Update:1 year ago
Show doc for...
Tokens:
Raw
### Pallas Quickstart

Source: https://docs.jax.dev/en/latest/_autosummary/jax.experimental.sparse.bcoo_todense

A quick start guide for using Pallas, demonstrating how to write and run a basic custom kernel.

```APIDOC
Pallas Quickstart:
  1. Define a kernel function using Pallas decorators and primitives.
  2. Specify input/output shapes and data types.
  3. Use `pallas.ops` for common operations like matrix multiplication or convolutions.
  4. Compile and run the kernel using JAX transformations.

  Example:
    import jax
    import jax.numpy as jnp
    import pallas as pl

    @pl.kernel
    def add_kernel(x_ref, y_ref, out_ref):
        val_x = x_ref[0]
        val_y = y_ref[0]
        out_ref[0] = val_x + val_y

    # Create a JAX array
    x = jnp.arange(10.0)
    y = jnp.arange(10.0) * 2
    out = jnp.zeros(10)

    # Execute the kernel
    pl.call_kernel(add_kernel, x, y, out)
    print(out)
```

--------------------------------

### Pallas Quickstart

Source: https://docs.jax.dev/en/latest/_autosummary/jax.scipy.stats.expon.pdf

A quick start guide for using Pallas, the JAX kernel language. This tutorial covers the basics of writing and running custom kernels with Pallas.

```APIDOC
pallas/quickstart.html
  - Description: Pallas Quickstart
  - Related: Pallas: a JAX kernel language, Software Pipelining, Grids and BlockSpecs
```

--------------------------------

### Install JAX CPU-only

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs the CPU-only version of JAX, suitable for local development on Linux, macOS, and Windows.

```bash
pip install -U jax
```

--------------------------------

### Detailed CPU Installation Steps

Source: https://docs.jax.dev/en/latest/_sources/installation

Provides detailed steps for CPU installation, including upgrading pip and handling Windows-specific requirements like the Visual Studio Redistributable.

```bash
pip install --upgrade pip
pip install --upgrade jax
```

--------------------------------

### Install JAX with TPU

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX with support for Google Cloud TPUs, primarily for use on TPU VMs.

```bash
pip install -U "jax[tpu]"
```

--------------------------------

### JAX vmap: Automatic Vectorization Example

Source: https://docs.jax.dev/en/latest/_sources/notebooks/thinking_in_jax

Demonstrates how jax.vmap can automatically vectorize a function, improving performance compared to naive Python loops or manual batching. It shows the setup for matrix-vector products and their transformation into matrix-matrix products.

```python
from jax import random
import jax.numpy as jnp
from jax import jit, vmap
import numpy as np

# Setup for matrix-vector product example
key = random.key(1701)
key1, key2 = random.split(key)
mat = random.normal(key1, (150, 100))
batched_x = random.normal(key2, (10, 100))

def apply_matrix(x):
  return jnp.dot(mat, x)

# Naively batched version using Python list comprehension
def naively_batched_apply_matrix(v_batched):
  return jnp.stack([apply_matrix(v) for v in v_batched])

# Manually batched version using jnp.dot's built-in semantics
@jit
def batched_apply_matrix(batched_x):
  return jnp.dot(batched_x, mat.T)

# Auto-vectorized version using jax.vmap
@jit
def vmap_batched_apply_matrix(batched_x):
  return vmap(apply_matrix)(batched_x)

# Assertions to verify correctness
np.testing.assert_allclose(naively_batched_apply_matrix(batched_x),
                           batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)
np.testing.assert_allclose(naively_batched_apply_matrix(batched_x),
                           vmap_batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)

print('Naively batched')
# %timeit naively_batched_apply_matrix(batched_x).block_until_ready()
print('Manually batched')
# %timeit batched_apply_matrix(batched_x).block_until_ready()
print('Auto-vectorized with vmap')
# %timeit vmap_batched_apply_matrix(batched_x).block_until_ready()
```

--------------------------------

### Install JAX (CPU/GPU)

Source: https://docs.jax.dev/en/latest/_sources/notebooks/thinking_in_jax

Installs JAX for CPU or NVIDIA GPU using pip. For GPU, specify the CUDA version. Refer to official documentation for detailed platform-specific instructions.

```shell
pip install jax
```

```shell
pip install -U "jax[cuda12]"
```

--------------------------------

### Install JAX Nightly (CPU)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs the latest JAX nightly build for CPU-only systems using pip. It specifies the index URL for nightly artifacts, ensuring the most recent development version is fetched.

```bash
pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
```

--------------------------------

### Detailed NVIDIA GPU Installation (CUDA via pip)

Source: https://docs.jax.dev/en/latest/_sources/installation

Details the recommended method for installing JAX with NVIDIA GPU support by using pip wheels for CUDA and cuDNN, simplifying the setup process.

```bash
pip install --upgrade pip

# NVIDIA CUDA 12 installation

```

--------------------------------

### Install Older JAXlib (CPU)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs a specific older version of the jaxlib wheel for CPU directly from the Python package index.

```bash
pip install jaxlib==0.3.25 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
```

--------------------------------

### Install Older JAX (CPU)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs a specific older version of JAX for CPU, using the Python package index. This command ensures a particular release version is installed.

```bash
pip install "jax[cpu]==0.3.25" -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
```

--------------------------------

### Install JAX with Google Cloud TPU Support

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX with support for Google Cloud TPUs. This command should be run within your Cloud TPU VM. It ensures the correct versions of jaxlib and libtpu are installed.

```bash
pip install "jax[tpu]"
```

--------------------------------

### Install JAX Nightly (Google Cloud TPU)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX nightly builds along with `libtpu` and `requests` for Google Cloud TPU support. It uses a specific index URL and a separate file URL for TPU releases.

```bash
pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

--------------------------------

### JAX: dynamic_slice_in_dim Example (1D)

Source: https://docs.jax.dev/en/latest/_autosummary/jax.lax.dynamic_slice_in_dim

Demonstrates the usage of jax.lax.dynamic_slice_in_dim with a one-dimensional array. It shows how to extract a slice of a specific size starting from a given index. The example also illustrates the behavior when the slice extends beyond the array bounds, showing that it gets clipped.

```python
import jax.numpy as jnp

x = jnp.arange(5)

# Slice from index 1 with size 3
print(jnp.dynamic_slice_in_dim(x, 1, 3))
# Expected output: Array([1, 2, 3], dtype=int32)

# Slice from index 4 with size 3 (out-of-bound, clipped)
print(jnp.dynamic_slice_in_dim(x, 4, 3))
# Expected output: Array([2, 3, 4], dtype=int32)
```

--------------------------------

### Pallas: A JAX Kernel Language Quickstart

Source: https://docs.jax.dev/en/latest/_autosummary/jax.custom_vjp.defvjp

Provides a quick start guide to Pallas, a domain-specific language (DSL) within JAX for writing high-performance custom kernels, particularly for accelerators like TPUs and GPUs.

```python
import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl

# Define a simple kernel using Pallas
@pallas.kernel
def add_kernel(x_ref, y_ref, o_ref):
  # Load values from input references
  x = x_ref[...] # Load all elements
  y = y_ref[...] # Load all elements
  # Perform computation
  o = x + y
  # Store result to output reference
  o_ref[...] = o

# Define the grid for execution (e.g., number of threads/work-items)
# This example assumes a simple 1D grid
grid_dim = jax.lax.create_token()

# Compile and run the kernel
# The kernel is applied to input/output arrays
input_x = jnp.ones(10)
input_y = jnp.ones(10)
output_arr = jnp.zeros(10)

# Execute the kernel with the defined grid
# The exact execution call depends on the Pallas API version and target hardware
# Example conceptual call:
# pl.run(add_kernel, grid_dim, input_x, input_y, output_arr)

# For a concrete example, refer to Pallas documentation for specific grid/block specs.

```

--------------------------------

### Install JAX from Source

Source: https://docs.jax.dev/en/latest/_sources/contributing

Steps to clone the JAX repository, install necessary testing requirements, and install JAX in editable mode for local development.

```bash
git clone https://github.com/YOUR_USERNAME/jax
cd jax
pip install -r build/test-requirements.txt
pip install -e ".[cpu]"
```

--------------------------------

### JAX Device Mesh Setup and Usage

Source: https://docs.jax.dev/en/latest/_sources/jep/9407-type-promotion

This snippet demonstrates how to set up a JAX device mesh and perform basic operations within it. It covers the creation of a mesh, defining logical axes, and distributing data across the mesh. The example uses JAX's pjit and related functions.

```python
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec, NamedSharding, PositionalSharding

# Define the device mesh
devices = jax.devices()
mesh_shape = (len(devices),)
mesh = Mesh(devices, ('data',))

# Define sharding specifications
pspec = PartitionSpec('data',)
named_sharding = NamedSharding(mesh, pspec)

# Example data
data = jnp.arange(len(devices))

# Apply sharding
sharded_data = jax.device_put(data, named_sharding)

print(f"Sharded Data: {sharded_data}")
```

--------------------------------

### Local Read the Docs Build Commands

Source: https://docs.jax.dev/en/latest/_sources/developer

Commands to set up a virtual environment, clone the JAX repository, checkout a specific branch, and build the documentation locally using Sphinx. Includes dependency installation and execution of the Sphinx build command.

```shell
mkvirtualenv jax-docs  # A new virtualenv
mkdir jax-docs  # A new directory
cd jax-docs
git clone --no-single-branch --depth 50 https://github.com/jax-ml/jax
cd jax
git checkout --force origin/test-docs
git clean -d -f -f
workon jax-docs

python -m pip install --upgrade --no-cache-dir pip
python -m pip install --upgrade --no-cache-dir -I Pygments==2.3.1 setuptools==41.0.1 docutils==0.14 mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 'sphinx<2' 'sphinx-rtd-theme<0.5' 'readthedocs-sphinx-ext<1.1'
python -m pip install --exists-action=w --no-cache-dir -r docs/requirements.txt
cd docs
python `which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html
```

--------------------------------

### JAX Tutorials Overview

Source: https://docs.jax.dev/en/latest/tutorials

This section lists available tutorials for learning JAX. These cover fundamental concepts such as just-in-time compilation, automatic differentiation, vectorization, and working with pytrees.

```APIDOC
JAX Tutorials:

  - Just-in-time compilation
  - Automatic vectorization
  - Automatic differentiation
  - Pseudorandom numbers
  - Stateful computations
  - Control flow and logical operators with JIT
  - Pytrees
  - Working with pytrees

For more in-depth discussions, refer to the advanced guides.
```

--------------------------------

### Install JAX with Local CUDA 12 Support

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX with support for a locally installed NVIDIA CUDA 12 toolkit. This is an alternative to using pre-built CUDA wheels. It requires a compatible NVIDIA driver and CUDA installation.

```bash
pip install --upgrade pip

# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 9.0 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12-local]"
```

--------------------------------

### JAX Sharding and Host Offloading Example

Source: https://docs.jax.dev/en/latest/notebooks/host-offloading

This Python code demonstrates setting up JAX sharding for device and pinned host memory. It shows how to move optimizer states to the host using `jax.device_put` to reduce device memory usage and analyzes the memory impact of this offloading strategy. The example includes initializing parameters and an optimizer, defining a step function, JIT compiling it with specific sharding, executing a step, and performing memory analysis.

```python
import jax
import jax.numpy as jnp
import optax

# Assume compute_loss and optimizer are defined elsewhere
# For demonstration, let's mock them:
def compute_loss(params, inputs):
    return jnp.sum(jnp.stack([jnp.sum(v) for v in params.values()]))

optimizer = optax.adam(learning_rate=0.1)

DIM = 128

# 1. Set up sharding specifications for device and host memory
s_dev = jax.sharding.SingleDeviceSharding(jax.devices()[0], memory_kind="device")
s_host = jax.sharding.SingleDeviceSharding(jax.devices()[0], memory_kind="pinned_host")

# Define the optimization step function
def step(params, opt_state, inputs):
  # Compute gradients
  grads = jax.grad(lambda p: compute_loss(p, inputs))(params)
  
  # Move optimizer state to the device for the update step
  # Note: In a real scenario, opt_state might already be on device or managed differently
  # This line shows explicit placement if it were on host and needed on device for update
  opt_state_on_device = jax.device_put(opt_state, s_dev)
  
  # Update parameters and optimizer state
  updates, new_opt_state = optimizer.update(grads, opt_state_on_device, params)
  new_params = optax.apply_updates(params, updates)
  
  # Return new parameters and the updated optimizer state (potentially still on device)
  return new_params, new_opt_state

# Initialize parameters
params = {f'w{i}': jnp.ones((DIM, DIM)) for i in range(1, 5)}

# Initialize optimizer state
# Optimizer state is placed on the host during initialization
initial_opt_state = optimizer.init(params)
opt_state_host = jax.device_put(initial_opt_state, s_host)

# JIT compile the step function with proper sharding and memory optimization
# donate_argnums=(0,) means 'params' can be mutated in place
# out_shardings specifies where the outputs should be placed
compiled_step = jax.jit(
  step,
  donate_argnums=(0,),
  out_shardings=(s_dev, s_host) # Output params on device, opt_state on host
)

# Mock input data
input_data = jnp.zeros((DIM, DIM))

# Run an optimization step
print("Running optimization step...")
new_params, offload_opt_state = compiled_step(params, opt_state_host, input_data)
print("Optimization step completed.")

# Analyze memory usage of the compiled function
# .lower() creates a version of the function that can be analyzed
# We need to provide concrete arguments to .lower()
print("Analyzing memory usage...")
compiled_step_for_analysis = compiled_step.lower(params, opt_state_host, input_data).compile()
compiled_stats = compiled_step_for_analysis.memory_analysis()

if compiled_stats is not None:
  # Calculate total memory usage as described in the text
  total_memory_bytes = (
      compiled_stats.temp_size_in_bytes 
      + compiled_stats.argument_size_in_bytes 
      + compiled_stats.output_size_in_bytes 
      - compiled_stats.alias_size_in_bytes
  )
  
  print(f"Memory Analysis Results:")
  print(f"  Temporary size: {compiled_stats.temp_size_in_bytes/(1024**3):.2f} GB")
  print(f"  Argument size: {compiled_stats.argument_size_in_bytes/(1024**2):.2f} MB") # Converted to MB for clarity as per text example
  print(f"  Total calculated size: {total_memory_bytes/(1024**3):.2f} GB")
else:
  print("Memory analysis not available for this compilation.")

# Note: The text mentions comparing numerical equivalence, which is omitted here for brevity.
# Example: jax.tree_util.tree_map(lambda p1, p2: jnp.allclose(p1, p2), params, new_params)

```

--------------------------------

### Install JAX with NVIDIA GPU (CUDA 12)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX with support for NVIDIA GPUs using CUDA 12, typically for Linux and macOS systems.

```bash
pip install -U "jax[cuda12]"
```

--------------------------------

### Install JAX (CPU)

Source: https://docs.jax.dev/en/latest/quickstart

Installs the JAX library for CPU-only execution on Linux, Windows, and macOS using pip.

```Python
pip install jax
```

--------------------------------

### JAX Installation Commands

Source: https://docs.jax.dev/en/latest/installation

Provides common pip installation commands for JAX, covering CPU-only, NVIDIA GPU with CUDA, and Google Cloud TPU configurations.

```python
pip install -U jax
```

```python
pip install -U "jax[cuda12]"
```

```python
pip install -U "jax[tpu]"
```

--------------------------------

### JAX Sharding Setup

Source: https://docs.jax.dev/en/latest/notebooks/explicit-sharding

This snippet demonstrates the basic setup for using JAX's sharding features, including importing necessary modules and configuring the number of CPU devices.

```python
import jax
import numpy as np
import jax.numpy as jnp
from jax.sharding import PartitionSpec as P, AxisType, get_abstract_mesh, reshard

jax.config.update('jax_num_cpu_devices', 8)
```

--------------------------------

### Install JAX with Conda

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs the JAX library using the conda package manager from the conda-forge channel. The second command specifically targets CUDA-enabled jaxlib packages for NVIDIA GPUs.

```bash
conda install jax -c conda-forge
```

```bash
conda install "jaxlib=*=*cuda*" jax -c conda-forge
```

--------------------------------

### Install Older JAXlib (GPU)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs a specific older version of the jaxlib wheel with CUDA support. It uses a dedicated URL for historical CUDA releases, specifying the CUDA and cuDNN versions.

```bash
pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

--------------------------------

### Initialize JAX Distributed System (GPU Example)

Source: https://docs.jax.dev/en/latest/_sources/multi_process

Demonstrates how to initialize JAX for distributed computing on GPU machines. It requires passing coordinator address, process ID, and total number of processes as command-line arguments to `jax.distributed.initialize`.

```python
# In file gpu_example.py...

import jax
import sys

# Get the coordinator_address, process_id, and num_processes from the command line.
coord_addr = sys.argv[1]
proc_id = int(sys.argv[2])
num_procs = int(sys.argv[3])

# Initialize the GPU machines.
jax.distributed.initialize(coordinator_address=coord_addr,
                           num_processes=num_procs,
                           process_id=proc_id)
print("process id =", jax.process_index())
print("global devices =", jax.devices())
print("local devices =", jax.local_devices())

```

--------------------------------

### JAX pjit Example

Source: https://docs.jax.dev/en/latest/jax.experimental.pjit

Demonstrates how to use `pjit` to compile a function with automatic partitioning over a mesh of devices. The example shows a convolution operation partitioned across available devices.

```python
import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, PartitionSpec
from jax.experimental.pjit import pjit

x = jnp.arange(8, dtype=jnp.float32)
f = pjit(lambda x: jax.numpy.convolve(x, jnp.asarray([0.5, 1.0, 0.5]), 'same'),
        in_shardings=None, out_shardings=PartitionSpec('devices'))

with Mesh(np.array(jax.devices()), ('devices',)):
  print(f(x))
```

--------------------------------

### Install JAX with CUDA 12 Support

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs JAX with support for NVIDIA CUDA 12. This command is primarily for Linux systems. Ensure your NVIDIA driver and CUDA toolkit versions are compatible.

```bash
pip install --upgrade "jax[cuda12]"
```

--------------------------------

### JAX vmap Usage Example Setup

Source: https://docs.jax.dev/en/latest/autodidax

Sets up a simple scalar function and input data to demonstrate the usage of the `vmap` transformation. This snippet prepares the environment for applying `vmap` to a basic operation.

```python
def add_one_to_a_scalar(scalar):
  assert np.ndim(scalar) == 0
  return 1 + scalar

vector_in = np.arange(3.)

```

--------------------------------

### JAX NamedSharding and Mesh Setup

Source: https://docs.jax.dev/en/latest/notebooks/host-offloading

Demonstrates setting up a JAX Mesh and NamedSharding for data distribution. It shows how to define a mesh representing device topology and prepare for sharding configurations, including memory kind specifications.

```python
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
import numpy as np

# Create mesh
# 1x1 mesh represents a single device with two named dimensions (x and y)
mesh = Mesh(np.array(jax.devices()[0]).reshape(1, 1), ('x', 'y'))

```

--------------------------------

### Install XProf TensorBoard Plugin

Source: https://docs.jax.dev/en/latest/profiling

Install the `xprof` package to get the TensorBoard Profiler plugin. Ensure only one version of TensorFlow or TensorBoard is installed to avoid conflicts. This enables visualization of JAX performance traces within TensorBoard.

```bash
pip install xprof
```

--------------------------------

### Install JAX Nightly (NVIDIA GPU CUDA 12 Legacy)

Source: https://docs.jax.dev/en/latest/_sources/installation

Installs legacy monolithic CUDA jaxlibs for older nightly releases. This option is generally not recommended as future monolithic CUDA jaxlibs will not be built.

```bash
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html
```

--------------------------------

### JAX Nightly Package Publication

Source: https://docs.jax.dev/en/latest/changelog

Informs users that JAX nightly packages are now published to artifact registry. Instructions for installing these packages can be found in the JAX installation guide.

```APIDOC
Installation:
  - JAX nightly packages are published to artifact registry.
  - Refer to the JAX installation guide for nightly installation instructions.
```

--------------------------------

### jax.example_libraries Overview

Source: https://docs.jax.dev/en/latest/jax.example_libraries

JAX includes experimental, small libraries designed as tools and examples for building machine learning libraries with JAX. These libraries are kept minimal (<300 source lines) to encourage adaptation and serve as inspiration rather than strict prescriptions.

```APIDOC
jax.example_libraries.optimizers
  Module providing basic optimization algorithms (e.g., SGD, Adam) as examples.
  Useful for understanding JAX-based optimization techniques.

jax.example_libraries.stax
  Module offering a simple neural network library built with JAX. It provides building blocks for defining and training neural networks, serving as a foundational example for more complex frameworks like Haiku or Flax.
```
Tokens:
Raw