Skip to content

← Flax step 7 of 100

Medium Primitives

Module That Branches on a Config Flag

Jump to the task

A configuration flag can choose whether a network includes an extra layer. In this Flax Linen exercise, the flag is a Python boolean stored on the module. It selects the architecture before the forward calculation runs.

Two configurations

Both configurations first map the input to hidden features and apply ReLU. When use_extra is true, add another Dense(hidden) followed by ReLU. Finish with a Dense(hidden) layer and no final activation.

Flag Sequence Output shape
False Dense, ReLU, Dense (hidden,)
True Dense, ReLU, Dense, ReLU, Dense (hidden,)

The final layer uses features=x.shape[-1] after the earlier layers have changed x. At that point the feature count is hidden, not the original input length. For a three-feature input and hidden=4, both configurations return four values. Negative output values are allowed.

Initialization and layer names

Define the layers inside @nn.compact. Initialization creates parameters only for the layers reached by the selected branch. With automatic names:

  • Dense_0 is the first layer in both configurations.
  • With the flag false, Dense_1 is the final layer.
  • With the flag true, Dense_1 is the extra layer and Dense_2 is the final one.

For input length D, the first kernel has shape (D, hidden). Every later kernel has shape (hidden, hidden), and all biases have shape (hidden,). The full variables mapping contains these arrays under "params".

Keep the same configuration for init and apply. A variables tree from the smaller model lacks Dense_2, which the larger model needs. In the other direction, a matching name can refer to a different role, so matching shapes alone do not make the trees interchangeable.

Explicit names can stabilize layer names in either setup or @nn.compact. For this exercise, use automatic names: names affect random initialization, and the expected values use the sequence above.

Configuration versus input-dependent branching

if self.use_extra: works here because the flag is a concrete Python boolean. Under jax.jit, an expression such as if x[0] > 0: tries to convert a traced input value to a Python boolean and raises TracerBoolConversionError. The same expression can work with a concrete scalar during eager execution; input arrays are not always tracers.

This exercise selects an architecture from configuration, not from the input values. JAX control-flow operations handle other kinds of data-dependent branching, but they are not needed here.

Problem

Implement conditional_layer_forward(seed, x, use_extra_layer, hidden):

  1. Convert the flag with bool(use_extra_layer >= 0.5) and store it as use_extra on ConditionalModel. Cast seed and hidden to int.
  2. Build the selected sequence above using @nn.compact, default Dense settings and unnamed layers. The final Dense reads the current x.shape[-1].
  3. Initialize the module with jax.random.PRNGKey(seed) and x, without splitting the key first.
  4. Pass the whole variables mapping to model.apply with the same x, and return the resulting JAX array of shape (hidden,).

x is a finite float32 array of shape (D,), with D >= 1. The seed is nonnegative, the hidden width is positive, and both are integer-valued scalars that may arrive as floats. The flag is 0.0 or 1.0. Leave x unchanged. The scalar conversions belong in this eager wrapper; it is not required to accept traced scalar arguments under jit.