JAX: Is This Google's NumPy killer?

jaxnumpygpu-accelerationautodiffpython

Abstraction: JAX library combining NumPy API with autodiff and GPU/TPU acceleration

Key points:

JAX provides jax.numpy (drop-in NumPy API), jax.grad (automatic differentiation), jax.jit (JIT compilation via XLA), jax.vmap (vectorization), jax.pmap (parallelization)
JIT compilation achieves ~100x speedup over NumPy: SELU on 10k×10k array runs in 0.357s (NumPy) vs 0.0024s (JAX JIT second run)
Key differences from NumPy: arrays are immutable; asynchronous execution requires block_until_ready(); random number generation requires explicit key management
JAX uses XLA compiler to target CPU, GPU, or TPU without code changes; first JIT call includes compilation overhead
Benchmarked examples: SELU activation function, automatic differentiation (cubic sum gradient), batched matrix-vector products, Gaussian image convolution
JAX is a Google Research project (not an official product), described as having "sharp edges"

Connections: Google · Jax · Numpy · Automatic Differentiation · Just In Time Compilation · Hardware Accelerators