Skip to content

← Flax NNX step 8 of 100

Medium Primitives

NNX Conditional Submodule

Jump to the task

Why this matters

A configuration flag can choose a model's architecture. This example optionally inserts a hidden layer between two Linear layers. When the flag is false, the extra layer is neither constructed nor called.

This is an architecture choice made when creating the model, not a per-example decision based on input data. The constructor and forward method use the same flag so the forward path only accesses layers that exist.

Worked example

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

class MaybeExtra(nnx.Module):
    def __init__(self, in_features, hidden, use_extra, rngs):
        self.linear1 = nnx.Linear(in_features, hidden, rngs=rngs)
        self.use_extra = use_extra
        if use_extra:
            self.extra = nnx.Linear(hidden, hidden, rngs=rngs)
        self.linear2 = nnx.Linear(hidden, hidden, rngs=rngs)

    def __call__(self, x):
        x = jax.nn.relu(self.linear1(x))
        if self.use_extra:
            x = jax.nn.relu(self.extra(x))
        return self.linear2(x)

without_extra = MaybeExtra(3, 4, False, rngs=nnx.Rngs(0))
with_extra = MaybeExtra(3, 4, True, rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)

assert not hasattr(without_extra, "extra")
assert hasattr(with_extra, "extra")
assert set(nnx.state(without_extra, nnx.Param)) == {"linear1", "linear2"}
assert set(nnx.state(with_extra, nnx.Param)) == {"linear1", "extra", "linear2"}
assert without_extra(x).shape == with_extra(x).shape == (4,)

Both models produce vectors of the same width, but their state structures differ. With the same seed, linear1 starts with the same parameters. Constructing extra advances the shared random stream before linear2 is initialized, so linear2 receives different keys in the two models.

The extra computation and changed initialization can both affect the output. Different architectures do not guarantee different values for every input: with zero input and these zero biases, both return zeros.

Static flags and traced values

The stored flag in this example is a Python bool in the model's graph definition. Split/merge preserves it. A Python if can use that known value, including when tracing a call through nnx.jit.

A dynamic traced Boolean is different. Python cannot choose a branch from a value that will only be known during execution, so converting it to bool raises a tracing error. It does not silently pick a valid branch. Constructors can also run during tracing if they are inside a compiled setup function; their inputs are not automatically static.

For a runtime decision, choose an appropriate JAX or NNX control-flow operation and compatible branches. jnp.where selects between array values; it is not a general replacement for conditionally constructing or calling a missing submodule. This exercise keeps the flag static.

Problem

Implement conditional_forward(seed, x, hidden, use_extra):

  1. Convert the configuration with use = bool(use_extra >= 0.5).
  2. Define MaybeExtra as shown above. Construct linear1, then extra only when enabled, then linear2. Pass the same Rngs instance to each.
  3. Use input width x.shape[-1], hidden width int(hidden) and nnx.Rngs(int(seed)), keeping the default Linear initializers and biases.
  4. Return the model's output on x. Apply ReLU after linear1 and after the optional extra layer, but not after linear2.

Input x is a nonempty finite float32 vector. Hidden is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; both may arrive as floats. The flag is supplied as 0.0 or 1.0. The specified conversion uses a 0.5 threshold, not ordinary nonzero truthiness. Return a float32 vector of length hidden.

Common pitfalls

  • Preserve construction order. Moving linear2 before extra changes its random keys and the expected output.
  • Keep allocation and use consistent. Changing a false flag to true after construction does not create the missing layer; build the intended model.
  • Do not interpret this example as a requirement to omit every unused layer. A different design may allocate both branches and select at runtime.
  • The final layer maps hidden to hidden in both configurations. The input width only affects linear1.
  • If JIT-compiling this setup function, mark seed, hidden and use_extra static. A changed static configuration can require a new compilation.

The numerical fixtures cannot prove that a layer was omitted from state. Use the example's attribute and state assertions to inspect that behavior.