nnx_variable___nnx_variablestate

Chapter 2: nnx.Variable / nnx.VariableState

In the previous chapter, we introduced nnx.Module as the foundation for building stateful neural network components in Flax NNX. We saw that modules hold their state (like parameters) as attributes. This chapter delves into the core mechanism NNX uses to represent and manage this state: nnx.Variable and its immutable counterpart, nnx.VariableState.

Motivation: Bridging Stateful Objects and Functional JAX

JAX thrives on pure functions operating on immutable data structures (pytrees). However, defining complex models often feels more natural using stateful object-oriented programming. How can we reconcile these two paradigms?

NNX uses nnx.Variable to encapsulate state within the familiar, mutable nnx.Module object. When it's time to interact with JAX transformations (jit, grad, etc.), NNX extracts this state into an immutable pytree structure where the leaves are nnx.VariableState objects. This separation allows developers to work with intuitive stateful objects while ensuring compatibility with JAX's functional nature.

• nnx.Variable: Lives inside the nnx.Module instance. It's a mutable container holding the actual state value (e.g., a JAX array) and associated metadata. You interact with it directly when defining or updating the module's state.
• nnx.VariableState: Exists outside the module instance, typically as part of the State pytree returned by nnx.split or nnx.state. It's an immutable snapshot of a Variable's value and metadata, suitable for use in JAX-transformed functions.

Central Use Case: Defining and Managing Different State Types

Let's enhance our nnx.Module with different kinds of state using various nnx.Variable subclasses.

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

# Define needed RNGs
rngs = nnx.Rngs(0)

class StatefulComponent(nnx.Module):
  def __init__(self, features: int, *, rngs: nnx.Rngs):
    # Learnable parameters
    self.weight = nnx.Param(
        initializers.lecun_normal()(rngs.params(), (features, features))
    )
    self.bias = nnx.Param(initializers.zeros_init()(rngs.params(), (features,)))
    
    # Batch statistics (e.g., for BatchNorm)
    self.running_mean = nnx.BatchStat(jnp.zeros((features,)))
    
    # RNG state (e.g., for Dropout)
    self.dropout_rng = nnx.RngState(rngs.dropout()) 
    
    # Custom state variable
    class StepCount(nnx.Variable): pass 
    self.steps = StepCount(0)

  def __call__(self, x: jax.Array):
    # Access variable values using .value
    output = jnp.dot(x, self.weight.value) + self.bias.value
    
    # Simulate using batch stats (simplified)
    # In reality, this would involve updating running_mean based on batch data
    output = output - self.running_mean.value 
    
    # Simulate using RNG state
    key = self.dropout_rng.value # Get current key
    key, subkey = jax.random.split(key)
    # In nnx.Dropout, the RngState variable is updated automatically
    # Here we simulate a manual update (not typical):
    # self.dropout_rng.value = key 
    noise = jax.random.normal(subkey, output.shape)
    output = output + noise
    
    # Update custom state
    self.steps.value += 1
    
    return output

  def update_mean(self, new_mean: jax.Array):
    # Directly mutate the Variable's value
    self.running_mean.value = new_mean

# Instantiate
component = StatefulComponent(features=4, rngs=rngs)

# Create dummy input
x = jnp.ones((1, 4))

# Call the component (modifies internal state)
y = component(x)
print(f"Initial step count: {component.steps.value}")
y = component(x)
print(f"Step count after second call: {component.steps.value}")

# Update a variable directly
component.update_mean(jnp.ones(4) * 0.1)
print(f"Updated running mean: {component.running_mean.value}")

# --- Interaction with Functional API ---

# Split the component into structure (GraphDef) and state (State)
graphdef, state = nnx.split(component)

print("\nExtracted State (containing VariableState objects):")
# Note: Output structure matches module attributes. Leaves are VariableState.
print(state) 

# The 'state' object is an immutable pytree suitable for JAX functions
# We can inspect the type and value of a specific state leaf
print(f"\nWeight state type: {type(state['weight'])}")
print(f"Weight value shape: {state['weight'].value.shape}")

# Reconstruct the component using nnx.merge
reconstructed_component = nnx.merge(graphdef, state)
print(f"\nReconstructed step count: {reconstructed_component.steps.value}")
assert component.steps.value == reconstructed_component.steps.value