JAX

JAX combines composable function transformations with automatic differentiation (Autograd) and Accelerated Linear Algebra (XLA) compilation for high-throughput machine learning research.

When to Use

• High-Performance Numerical & Scientific Computing: Autograd and XLA compilation for accelerated linear algebra on GPU/TPU.
• Custom Deep Learning Research: Building neural networks with Flax, Haiku, or Equinox without framework bloat.
• Composable Function Transformations: Seamlessly chaining jit (compilation), grad (derivatives), and vmap (vectorization).
• Multi-Device Distributed Training: Parallelizing computation across clusters using pmap and shard maps (jax.experimental.shard_map).

Quick Start

import jax
import jax.numpy as jnp