Skip to content

← Flax NNX step 18 of 100

Medium Primitives

NNX Plain Array Attributes

Jump to the task

Why this matters

A model may use an array without marking it as a Param. For example, a precomputed table can participate in the forward calculation while staying outside a parameter-only update.

Here, the kernel is wrapped in Param and the bias is a plain JAX array. A Param filter selects only the kernel. This does not make the bias immutable, static or incapable of receiving gradients. It describes what that particular filter selects.

Worked example

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

class FrozenAttr(nnx.Module):
    def __init__(self, in_features, out_features, frozen_bias_value, rngs):
        key = rngs.params()
        self.kernel = nnx.Param(
            jax.random.normal(key, (in_features, out_features))
            * (1.0 / jnp.sqrt(in_features))
        )
        self.frozen_bias = jnp.ones((out_features,), dtype=jnp.float32) * frozen_bias_value

    def __call__(self, x):
        return x @ self.kernel + self.frozen_bias

model = FrozenAttr(3, 4, 2.0, rngs=nnx.Rngs(0))
assert set(nnx.state(model, nnx.Param)) == {"kernel"}
assert set(nnx.state(model, nnx.Variable)) == {"kernel"}
assert set(nnx.state(model)) == {"kernel", "frozen_bias"}
assert len(jax.tree.leaves(nnx.state(model, nnx.Param))) == 1
assert jnp.linalg.norm(model.frozen_bias) == 4.0

The existing class and argument names use frozen for compatibility with the exercise. The bias is only excluded from the Param selection. Unfiltered state includes it, and split/merge can carry it along with the kernel.

Filtering is not freezing

A type filter does not tell Optax what to do on its own. The code that supplies values and gradients to an optimizer determines what it updates. NNX gradient and optimizer filters can select parameters, while other differentiation paths can still involve a plain array.

For example, differentiating sum(bias) with respect to bias gives a vector of ones. Leaving bias unwrapped does not insert stop_gradient. You can also replace model.frozen_bias with another array in ordinary Python. A dataclass field with init=False would only affect constructor arguments; it would not enforce immutability.

Bare arrays are not automatically static metadata. In the supported NNX versions, this array participates in model state and can be passed through nnx.jit. Use wrappers when you need a named variable role, not because every unwrapped array is excluded from transformations.

Problem

Implement frozen_kernel_unchanged(seed, x, features, frozen_bias_value):

  1. Define FrozenAttr as shown above. Consume one params key for the normal kernel, scaled by 1 / sqrt(in_features).
  2. Construct it with input width x.shape[-1], output width int(features), the supplied bias value and nnx.Rngs(int(seed)).
  3. Count array leaves selected by nnx.state(model, nnx.Param).
  4. Compute the kernel's Frobenius norm and the bias vector's Euclidean norm.
  5. Return a float32 array [leaf_count, kernel_norm, bias_norm].

Input x is a nonempty float32 vector. The function uses only its width; it does not need to call the forward method. Features is a positive integer-valued scalar, seed is a nonnegative integer-valued scalar, and both may arrive as floats. The bias value is a finite scalar.

The return shape is (3,). Its first entry is 1, and its final entry is abs(frozen_bias_value) * sqrt(features). No optimizer step occurs in this exercise, despite the historical function name.

Common pitfalls

  • Keep the bias as a plain array to practice this filter behavior. A Param wrapper would make it a second leaf in the parameter selection.
  • Read the kernel with model.kernel[...] and the plain bias directly.
  • Count array leaves, not the kernel's numeric entries.
  • Keep the bias value and both norms as JAX values. Stack the output scalars instead of converting them to Python floats.
  • If applying JIT to this setup function, mark seed and features static. The bias value can remain a dynamic argument.

The numerical fixtures cannot establish how arrays were wrapped or filtered. Use the example's state assertions to inspect that structure.