jax

JAX - Autograd and XLA (Accelerated Linear Algebra)

JAX is a framework that combines a NumPy-like API with a powerful system of composable function transformations: Grad (differentiation), Jit (compilation), Vmap (vectorization), and Pmap (parallelization).

When to Use

• High-performance scientific simulations requiring GPU/TPU acceleration.
• Custom machine learning research where PyTorch/TF abstractions are too restrictive.
• Calculating higher-order derivatives (Hessians, Jacobians) for optimization.
• Physics-informed machine learning and differentiable simulations.
• Automatic vectorization of functions (no more manual batching).
• Running the same code on CPU, GPU, and TPU without changes.

Reference Documentation

Official docs: https://jax.readthedocs.io/
GitHub: https://github.com/google/jax
Search patterns: jax.numpy, jax.jit, jax.grad, jax.vmap, jax.random

Core Principles

Pure Functions (Immutability)

JAX is built on functional programming. All functions must be pure: they should not have side effects (like modifying a global variable) and must return the same output for the same input. JAX arrays are immutable.

XLA (Just-In-Time Compilation)

JAX uses XLA to compile and optimize Python/NumPy code into efficient machine code for specific hardware.

Manual PRNG Handling

Unlike NumPy, JAX requires explicit management of random state (keys) to ensure reproducibility in parallel environments.

Quick Reference

Installation

# CPU
pip install jax jaxlib

# GPU (Check documentation for specific CUDA versions)
pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Standard Imports

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap, pmap, random

Basic Pattern - Differentiate and JIT

import jax.numpy as jnp
from jax import grad, jit

# 1. Define a pure function
def f(x):
    return jnp.sin(x) + x**2

# 2. Transform: Create a gradient function
df_dx = grad(f)

# 3. Transform: Compile for speed
f_fast = jit(f)

# 4. Use
val = f_fast(2.0)
slope = df_dx(2.0)