Flax

Warnings

• breaking The `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.
• gotcha Flax, 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)`.
• gotcha JAX'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.
• gotcha Flax `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.

Install

• pip install flax jax[cpu] CPU-only JAX
• pip install flax jax[cuda12_pip] CUDA JAX (adjust for your CUDA version)

Imports

• nn

  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.

Quickstart

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})