filters___filterlib__

Chapter 4: Filters (filterlib)

In the previous chapters, we learned how nnx.Module defines network components, how nnx.Variable / nnx.VariableState represents their state, and how nnx.Rngs manages randomness. Models often contain diverse state elements (parameters, batch statistics, caches, RNG state, etc.) nested within complex structures. When working with these models, especially when using the NNX Functional API (split/merge/state/update/graphdef) or applying NNX Lifted Transforms (jit, grad, vmap, scan, etc.), we frequently need to select specific subsets of this state. This is where Filters come in.

Motivation: Targeted State Manipulation

Imagine you want to:

• Apply an optimizer update only to the learnable parameters (nnx.Param) and not the batch statistics (nnx.BatchStat).
• Extract only the batch statistics to update them manually during inference.
• Specify which parts of the state should be differentiated using nnx.grad.
• Define how different state components should be handled (mapped, broadcasted, carried) in nnx.vmap or nnx.scan.
• Remove temporary state like intermediate activations (nnx.Intermediate) after a forward pass.

Manually traversing the nested State pytree and checking types or names for each VariableState would be verbose and error-prone. Filters provide a concise and powerful declarative language for specifying these selections.

Central Use Case: Separating Parameters and Batch Statistics

A common task is separating the parameters that need gradients from the batch statistics that are updated differently. Filters make this straightforward using nnx.split.

import jax
import jax.numpy as jnp
from flax import nnx

class MyModel(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.layer1 = nnx.Linear(10, 20, rngs=rngs)
    self.bn = nnx.BatchNorm(20, use_running_average=False, rngs=rngs)
    self.layer2 = nnx.Linear(20, 5, rngs=rngs)

  def __call__(self, x):
    x = self.layer1(x)
    x = self.bn(x)
    x = nnx.relu(x)
    x = self.layer2(x)
    return x

model = MyModel(rngs=nnx.Rngs(0))

# Use Variable types as filters in nnx.split
graphdef, params_state, bn_state, other_state = nnx.split(
    model, 
    nnx.Param,       # Filter 1: Select all nnx.Param variables
    nnx.BatchStat,   # Filter 2: Select all nnx.BatchStat variables
    ...              # Filter 3: Select everything else (Ellipsis or True)
)

print("--- Parameters (nnx.Param) ---")
# params_state contains only VariableState objects whose type is nnx.Param
print(params_state) 

print("\n--- Batch Statistics (nnx.BatchStat) ---")
# bn_state contains only VariableState objects whose type is nnx.BatchStat
print(bn_state)

print("\n--- Other State (e.g., RngState for BatchNorm) ---")
# other_state contains the rest (e.g., RngState from BatchNorm)
print(other_state)

# Reconstruct the model
model_merged = nnx.merge(graphdef, params_state, bn_state, other_state)