flax

library0.12.9pypypi✓ verified 30d ago

Flax is a high-performance neural network library for JAX, designed for flexibility and ease of use. It provides building blocks for defining models, handling parameters, and managing training state within the JAX ecosystem. As of my last check, the current version is 0.12.6. Its release cadence is closely tied to JAX updates and major developments in the JAX ecosystem, with frequent minor and patch releases.

ai-ml

pip install flax jax[cpu]

github ↗package ↗

INSTALL

IMPORT

SIG · FLAX

flax

ai-mlpythonv0.12.9

Install

20.3s avg

Import

3301ms

Disk

650MB

Pass rate

6/ 10

Env Coverage6 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.3.3 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

1/2 runs

✓ 18.45s

py 3.11

✓ 0.1s

✓ 18.95s

py 3.12

✕ timeout

✓ 17.1s

py 3.13

✕ timeout

✓ 17.25s

py 3.9

1/2 runs

✓ 49.9s

650MB installed

● package 650MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

✓ import flax.linen as nn

The primary module for defining neural network layers and modules.

FrozenDict

✓ from flax.core import FrozenDict

Used for immutable parameter structures returned by model initialization.

TrainState

✓ from flax.training import train_state

Common utility for managing model parameters, optimizer state, and other training state.

This quickstart demonstrates how to define a basic neural network module using `flax.linen`, initialize its parameters with a JAX PRNG key, and perform a forward pass. It also shows how to pass explicit PRNG keys for stochastic operations.

import jax import jax.numpy as jnp import flax.linen as nn # Define a simple Multi-Layer Perceptron (MLP) class MLP(nn.Module): num_neurons: int @nn.compact def __call__(self, x): x = nn.Dense(features=self.num_neurons)(x) # First dense layer x = nn.relu(x) x = nn.Dense(features=self.num_neurons)(x) # Second dense layer return x # Example usage: key = jax.random.PRNGKey(0) # Initialize a PRNG key model = MLP(num_neurons=64) # Create a dummy input (batch_size, input_features) dummy_input = jnp.ones((1, 10)) # Initialize model parameters # The 'params' are stored in a FrozenDict within the initialized variables variables = model.init(key, dummy_input) params = variables['params'] print(f"Initial parameters structure: {jax.tree_map(lambda x: x.shape, params)}") # Perform a forward pass output = model.apply({'params': params}, dummy_input) print(f"Output shape: {output.shape}") # Example of applying with a different PRNG key for randomness (e.g., dropout) dropout_key, _ = jax.random.split(key) output_with_rng = model.apply({'params': params}, dummy_input, rngs={'dropout': dropout_key})

Debug

Known issues

breakingThe `flax.optim` module, which previously provided optimizers, has been deprecated since Flax 0.5.0 and fully removed in later versions. Attempting to use it will result in import errors.

fix

Migrate to `optax` for all optimizer definitions and applications. `flax.training.train_state` is designed to work seamlessly with `optax` optimizers.

affects: >=0.5.0

gotchaFlax, built on JAX, enforces immutability for all parameters and model states. Operations that modify state (e.g., weight updates during training) do not mutate in place but return a *new* Pytree with the updated values. This is crucial for JAX's functional paradigm.

fix

Always assign the result of any state-modifying function (e.g., `optimizer.apply_gradients` or `state.apply_gradients`) back to your state variable: `state = state.apply_gradients(grads)`.

affects: All versions

gotchaJAX's functional approach requires explicit handling and splitting of pseudo-random number generator (PRNG) keys for any stochastic operation (e.g., `Dropout`, `initializers`). Reusing the same key will produce the same 'random' sequence, and failing to split keys can lead to non-random or deterministic behavior where randomness is expected.

fix

For operations requiring randomness within `nn.Module` (e.g., dropout), pass a dictionary of PRNG keys to `model.apply()` using the `rngs` argument (e.g., `model.apply(..., rngs={'dropout': dropout_key})`). Use `jax.random.split` to generate new sub-keys for subsequent random operations.

affects: All versions

gotchaFlax `nn.Module`s typically initialize parameters (and other variables) based on input shapes during `model.init()`. If input shapes change during inference or subsequent calls, the model's structure or behavior might implicitly change or lead to errors if not handled correctly.

fix

Ensure that the input shape used for `model.init()` accurately reflects the expected input shape during the model's lifetime. For dynamic shapes, consider using `jax.ShapeDtypeStruct` or designing modules to be robust to varying batch sizes or sequence lengths.

affects: All versions

Upgrade

Version history

0.12.9latest on PyPI · released Aug 18, 2026

Audit

Dependencies

jaxrequiredFlax is built on top of JAX and requires it for all operations.

optaxoptionalThe recommended library for optimizers in Flax, replacing the deprecated flax.optim.

Agent activity

14 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/google/flax ↗

packagehttps://pypi.org/project/flax/ ↗