Why JAX Excels at Scientific Computing

JAX was born from the scientific computing community, and its core features — automatic differentiation, JIT compilation, and vectorization — are exactly what scientists need. While PyTorch and TensorFlow were designed primarily for neural networks, JAX was designed for differentiable computation in general. This makes it uniquely powerful for scientific applications.

Consider what a scientist might need:

• Gradients of any computation — not just neural network loss functions, but physical simulations, statistical models, optimization objectives
• Fast numerical code — JIT compilation gives you C++/Fortran-level speed from Python
• Vectorization — run the same simulation with 1000 different initial conditions, simultaneously
• Hardware acceleration — GPUs and TPUs for massive parallelism

import jax
import jax.numpy as jnp

# Example: differentiable physics simulation
def simulate_spring(k, m, x0, v0, dt, num_steps):
    """Simulate a damped spring system: m*x'' + 0.1*x' + k*x = 0"""
    def step(state, _):
        x, v = state
        a = (-k * x - 0.1 * v) / m  # spring force + damping
        v_new = v + a * dt
        x_new = x + v_new * dt
        return (x_new, v_new), x_new

    init_state = (x0, v0)
    _, trajectory = jax.lax.scan(step, init_state, None, length=num_steps)
    return trajectory

# Simulate
traj = simulate_spring(k=2.0, m=1.0, x0=1.0, v0=0.0, dt=0.01, num_steps=1000)

# Gradient: how does the final position change with spring constant?
@jax.jit
def final_position(k):
    return simulate_spring(k, 1.0, 1.0, 0.0, 0.01, 1000)[-1]

dk = jax.grad(final_position)(2.0)
print(f"d(final_pos)/dk = {dk:.6f}")

# Vectorize: simulate 100 different spring constants at once
ks = jnp.linspace(0.5, 5.0, 100)
all_trajectories = jax.vmap(lambda k: simulate_spring(k, 1.0, 1.0, 0.0, 0.01, 1000))(ks)
print(f"Batch trajectories shape: {all_trajectories.shape}")  # (100, 1000)