NumPy to JAX: The Array That Learned New Tricks

NumPy stores numbers in arrays whose operations run in compiled code. JAX keeps that array API and adds compilation, automatic differentiation, and accelerator backends.

1 History

Python had two earlier array libraries:

Numeric — fast on small arrays.
Numarray — features needed for large astronomical images.

The two APIs were incompatible. Travis Oliphant merged them and released NumPy 1.0 in 2006 (Guide to NumPy).

That common type is what pandas, SciPy, and scikit-learn exchange. Harris et al. (2020), “Array programming with NumPy”, traces the downstream impact.

2 NumPy primitives

Four constructors and operations cover most work:

np.array — Python values to an array.
np.arange — regular sequence.
.reshape() — change dimensions without changing values.
elementwise arithmetic — no Python loop.

Array metadata:

.shape — length of each axis.
.ndim — number of axes.
.size — total element count.
.dtype — storage type.

The new dimensions after reshape must still multiply to .size.

Indexing:

grid[1, 2] — one value.
grid[1, :] — one row.
grid[:, 2] — one column.

Vectorized math: state the array operation; NumPy runs it in compiled routines.

3 Reductions and broadcasting

Reductions collapse many values to summaries: .mean(), .sum(), .min(), .max().

On a matrix, axis picks the direction:

grid.sum(axis=0) — column totals.
grid.sum(axis=1) — row totals.

Broadcasting subtracts a scalar mean from every element.

The check should be near zero. Residual is floating-point rounding.

These values are synthetic: five numbers chosen for a centering demo, not measured.

4 JAX

JAX mirrors the NumPy API:

import jax.numpy as jnp

x = jnp.arange(12).reshape(3, 4)
y = jnp.sin(x) + x**2

Additions relative to NumPy:

GPU / TPU execution for compatible work.
jax.jit — compile a function for repeated runs.
jax.grad — return a function that computes derivatives.

Frostig, Johnson, and Leary (2018), “Compiling machine learning programs via high-level tracing”, describes the tracing used for compilation.

Characteristic pattern:

import jax
import jax.numpy as jnp

def loss(weight):
    return (weight - 3.0) ** 2

fast_loss = jax.jit(loss)
loss_slope = jax.grad(loss)

print(fast_loss(1.0))   # 4.0
print(loss_slope(1.0))  # -4.0

jaxlib has no Pyodide build, so JAX cannot run in this page. The live cell below uses a finite-difference slope on the same scalar loss.

jax.grad(loss)(weight) returns the exact slope without a finite step.

5 Constraints

JAX transformations assume pure functions: output depends only on inputs; no hidden mutation of outside state.

Arrays are immutable. Replace x[0] = 10 with:

x = x.at[0].set(10)

Those rules make compilation and differentiation well-defined.

6 Choice of library

Use NumPy for general analysis on CPU.
Use JAX when gradients, jit, or accelerators dominate the workload.

Learn NumPy first. Its array model is the shared grammar; JAX reuses it.

7 References

Oliphant, T. E. Guide to NumPy — Numeric/Numarray merger and NumPy 1.0 (2006).
Harris, C. R. et al. (2020). Array programming with NumPy. Nature.
Frostig, R., Johnson, M. J. & Leary, C. (2018). Compiling machine learning programs via high-level tracing. MLSys.