If you come from PyTorch, you're used to torch.nn being part of the core framework. TensorFlow has tf.keras. But JAX intentionally ships without a built-in neural network library. This is a deliberate design choice, not an oversight.
JAX's philosophy is separation of concerns: JAX provides the mathematical primitives (arrays, autodiff, compilation, parallelism) and lets the community build neural network abstractions on top. This has led to several excellent libraries, each with different design philosophies:
- Flax (Google) — the "official" neural network library, now using the NNX API
- Equinox (Patrick Kidger) — models as pytrees, PyTorch-like feel
- Haiku (DeepMind) — now in maintenance mode, legacy codebases
This ecosystem approach has trade-offs:
# PyTorch: one way to define a model
# class Model(nn.Module):
# def __init__(self):
# super().__init__()
# self.linear = nn.Linear(784, 10)
# def forward(self, x):
# return self.linear(x)
# JAX: you choose your library
# Flax NNX version:
from flax import nnx
class Model(nnx.Module):
def __init__(self, rngs):
self.linear = nnx.Linear(784, 10, rngs=rngs)
def __call__(self, x):
return self.linear(x)
# Equinox version:
import equinox as eqx
class Model(eqx.Module):
linear: eqx.nn.Linear
def __init__(self, key):
self.linear = eqx.nn.Linear(784, 10, key=key)
def __call__(self, x):
return self.linear(x)💡 Why This Matters
The JAX ecosystem's diversity is both its greatest strength and its biggest hurdle for newcomers. But the payoff is real: each library can innovate independently, and they all share the same JAX foundation. Flax NNX is the default choice for most projects (especially in the Google ecosystem), while Equinox is popular in scientific ML. You don't need to learn all of them — pick one and you can always switch later, since the underlying JAX concepts are the same.
The JAX AI Stack (pip install jax-ai-stack) bundles the most common libraries together: JAX + Flax + Optax + Orbax + ml_dtypes. This is the easiest way to get a complete deep learning setup with JAX.