repository·main·Indexed 27 days ago
https://github.com/google/flaxA high-performance neural network library and ecosystem for JAX designed for flexibility. Flax provides a neural network API, including the Pythonic NNX and the stable Linen API, along with utilities for replicated training, serialization, checkpointing, and metrics.
Flax NNX is a neural network API designed to bring Pythonic reference semantics to JAX. Unlike the functional (lazy initialization) semantics of Flax Linen, NNX uses regular Python semantics for nnx.Module, supporting mutability and shared references.
Key features include:
jax.vmap, jax.grad, or jax.scan) that are specifically designed to be applied to Flax Modules. Instead of applying transformations directly to raw functions, you use the lifted versions provided by Flax to ensure module state and RNGs are handled correctly.Linen is the recommended neural network API for Flax (replacing the deprecated flax.nn API). It is designed to be a 'comfortable evolution' that improves submodule sharing and support for non-trainable variables.
Key features include:
vmap, remat, or scan directly inside your modules.Flax NNX is designed to improve upon Flax Linen by providing a more intuitive, object-oriented experience. Key improvements include:
nnx.Module objects are regular Python objects. Unlike Linen modules, which are lazy and difficult to inspect before runtime, NNX modules can be constructed and inspected immediately.apply or init. Parameters are held as attributes, and methods (including __call__ and custom methods) can be called directly on the module instance.BatchNorm statistics or Dropout flags) is kept inside the nnx.Module and is mutable. This eliminates the need to manually manage complex parameter/state dictionaries during training loops.nnx.vmap and other NNX transforms are designed to be equivalent to JAX transforms. They can be used anywhere (including training loops), can accept modules as any argument, and can return modules.Flax Linen is a functional system for defining neural networks in JAX. Unlike object-oriented frameworks, Linen is designed to be compatible with JAX's composable function-transformation approach.
Key architectural concepts include:
f(v_in, x) -> v_out, y, where v_in represents variable collections and PRNG state, x is input data, v_out is the mutated variable collections, and y is the output data.params, batch_norm, cache) to allow different treatment of variables under JAX transformations (like jit, vmap, or pmap).In NNX, you can perform stateful updates (like updating running averages in BatchNorm) by creating a Variable and updating its .value property during the forward pass. You can use in-place update syntax like self.variable[...] += 1.
class Count(nnx.Variable): pass
class Counter(nnx.Module):
def __init__(self):
self.count = Count(jnp.array(0))
def __call__(self):
self.count[...] += 1
counter = Counter()
counter()
print(f'{counter.count[...] = }')Flax Modules are mutable by default. You can modify the model structure at any time by replacing sub-Module attributes or Variables with new ones (e.g., replacing a standard layer with a LoRA layer).
# Example: Replacing linear layers with LoraLinear
model.linear1 = LoraLinear(model.linear1, 4, rngs=rngs)
model.linear2 = LoraLinear(model.linear2, 4, rngs=rngs)nnx.avg_pool() does not have a direct equivalent to PyTorch's count_include_pad=False (which excludes zero-padding from the average calculation).
To replicate this behavior, you can use nnx.pool() with jax.lax.add to manually calculate the average by dividing the sum by the count of non-padded elements.
def avg_pool(inputs, window_shape, strides=None, padding='VALID'):
"""
Pools the input by taking the average over a window.
Does not consider padded zero's for the average computation.
"""
assert len(window_shape) == 2
y = nnx.pool(inputs, 0., jax.lax.add, window_shape, strides, padding)
counts = nnx.pool(jnp.ones_like(inputs), 0., jax.lax.add, window_shape, strides, padding)
y = y / counts
return ySince Flax v0.6.0, flax.optim has been removed. You should migrate to Optax. Unlike flax.optim, optax does not keep a copy of the parameters (params). You must manage params and opt_state separately.
To simplify state management, Flax provides flax.training.train_state.TrainState to store optimizer state, parameters, and other associated data in a single dataclass.
@jax.jit
def train_step(params, opt_state, batch):
grads = jax.grad(loss)(params, batch)
updates, opt_state = tx.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return params, opt_state
tx = optax.sgd(learning_rate, momentum)
params = variables['params']
opt_state = tx.init(params)
for batch in ds_train:
params, opt_state = train_step(params, opt_state, batch)State (such as BatchNorm statistics or Dropout randomness) is stored directly within the nnx.Module and is mutable. This allows you to add stateful layers without changing your training loop logic.
class BatchNorm(nnx.Module):
def __init__(self, features: int, mu: float = 0.95):
self.scale = nnx.Param(jax.numpy.ones((features,)))
self.bias = nnx.Param(jax.numpy.zeros((features,)))
self.mean = nnx.BatchStat(jax.numpy.zeros((features,)))
self.var = nnx.BatchStat(jax.numpy.ones((features,)))
self.mu = mu
def __call__(self, x):
mean = jax.numpy.mean(x, axis=-1)
var = jax.numpy.var(x, axis=-1)
# ema updates happen directly on the stateful attributes
self.mean.value = self.mu * self.mean + (1 - self.mu) * mean
self.var.value = self.mu * self.var + (1 - self.mu) * var
x = (x - mean) / jax.numpy.sqrt(var + 1e-5)
return x * self.scale + self.bias