JAX and Automatic Differentiation¶

NLSQ uses JAX for GPU acceleration and automatic Jacobian computation. This guide explains how it works and why it matters for curve fitting.

What is JAX?¶

JAX is a numerical computing library that provides:

NumPy-compatible API: Write code like you would with NumPy
Automatic differentiation: Compute derivatives automatically
JIT compilation: Compile Python to optimized machine code
GPU/TPU acceleration: Run on accelerators with no code changes

Why JAX for Curve Fitting?¶

Traditional approach (SciPy):

Compute Jacobian using finite differences
Each partial derivative requires a function evaluation
For m parameters: 2m extra function calls per iteration
Numerically approximate (subject to step size errors)

JAX approach (NLSQ):

Compute exact Jacobian via automatic differentiation
Single backward pass computes all derivatives
No extra function evaluations
Analytically exact (machine precision)

# You write this
def model(x, a, b, c):
    return a * jnp.exp(-b * x) + c


# JAX automatically computes
# ∂f/∂a = exp(-b * x)
# ∂f/∂b = -a * x * exp(-b * x)
# ∂f/∂c = 1

Automatic Differentiation¶

AD is not: - Symbolic differentiation (like Mathematica) - Numerical differentiation (finite differences)

AD is: - Algorithmic transformation of code - Tracks derivatives through computations - Exact to machine precision

Two Modes¶

Forward mode: Propagate derivatives forward through computation

Input: x     →     f₁(x)     →     f₂(f₁(x))     →     Output
       ∂x/∂x=1 → ∂f₁/∂x     →     ∂f₂/∂x        →     ∂y/∂x

Good for few inputs, many outputs.

Reverse mode (backpropagation): Propagate derivatives backward

Input: x     →     f₁(x)     →     f₂(f₁(x))     →     Output
       ∂y/∂x ←     ∂y/∂f₁   ←     ∂y/∂f₂=1     ←     ∂y/∂y=1

Good for many inputs, few outputs (like gradients in optimization).

NLSQ uses reverse mode to efficiently compute the Jacobian.

JIT Compilation¶

JAX’s Just-In-Time compiler transforms Python to optimized XLA code:

@jax.jit
def model(x, a, b):
    return a * jnp.exp(-b * x)


# First call: compile (slower)
y1 = model(x, 1.0, 0.5)

# Subsequent calls: run compiled code (fast!)
y2 = model(x, 2.0, 0.3)

Benefits:

Operator fusion: Combine multiple operations
Memory optimization: Reduce intermediate allocations
Parallelization: Utilize all CPU cores or GPU threads
Constant folding: Pre-compute static values

Why Use jax.numpy?¶

JAX operations must be traced to enable AD and JIT:

import numpy as np
import jax.numpy as jnp


# This WON'T work with JAX
def bad_model(x, a, b):
    return a * np.exp(-b * x)  # NumPy exp can't be traced


# This WORKS with JAX
def good_model(x, a, b):
    return a * jnp.exp(-b * x)  # JAX exp is traceable

Rule: Use jax.numpy for any math inside model functions.

GPU Acceleration¶

JAX automatically uses GPU when available:

import jax

# Check available devices
print(jax.devices())  # [cuda(id=0)] or [cpu()]

# Data automatically moves to GPU
x = jnp.array([1, 2, 3])  # Lives on GPU if available

# Computations run on GPU
y = jnp.exp(x)  # Computed on GPU

No code changes needed - same code runs on CPU or GPU.

Pure Functions¶

JAX requires pure functions - no side effects:

# BAD: Side effects
counter = 0


def bad_model(x, a, b):
    global counter
    counter += 1  # Side effect!
    return a * jnp.exp(-b * x)


# GOOD: Pure function
def good_model(x, a, b):
    return a * jnp.exp(-b * x)  # No side effects

Why? JAX may: - Cache and reuse results - Execute operations in different order - Run computations in parallel

Common Gotchas¶

Dynamic shapes

# BAD: Shape depends on values
def bad(x, a):
    if a > 0:  # Python control flow on traced value
        return x[:10]
    return x


# GOOD: Use jnp.where for conditionals
def good(x, a):
    return jnp.where(a > 0, x * 2, x)

In-place mutation

# BAD: Mutating arrays
def bad(x):
    x[0] = 0  # JAX arrays are immutable!
    return x


# GOOD: Create new array
def good(x):
    return x.at[0].set(0)

Random numbers

# BAD: NumPy random
def bad():
    return np.random.randn()  # Not reproducible in JAX


# GOOD: JAX random with key
def good(key):
    return jax.random.normal(key)

Performance Tips¶

Warm up JIT

# First call compiles (slow)
_ = model(x_small, *p0)

# Subsequent calls are fast
result = model(x_large, *p0)

Batch similar computations

# Use vmap for vectorization
batched_model = jax.vmap(model, in_axes=(0, None, None))
results = batched_model(x_batch, a, b)

Use streaming optimizer for larger datasets

from nlsq import curve_fit_large

popt, pcov = curve_fit_large(model, x, y, p0=p0)  # Memory-efficient

Summary¶

JAX enables NLSQ’s key features:

Automatic Jacobians: Exact derivatives, no manual math
JIT compilation: Fast execution after first call
GPU acceleration: Same code, massive speedups
Numerical precision: IEEE 754 exact derivatives

Just remember to use jax.numpy in your model functions!