Skip to content

← Flax step 6 of 100

Medium Primitives

Module with Multiple Named Sub-Modules

Jump to the task

Build a module from two named Dense children: an encoder that maps the input to a hidden vector and a decoder that maps it back to the input width. This exercise practices module composition and parameter paths. It does not train the model to reconstruct its input.

Define the children in setup

class EncoderDecoder(nn.Module):
    hidden: int
    out: int

    def setup(self):
        self.encoder = nn.Dense(features=self.hidden, name="encoder")
        self.decoder = nn.Dense(features=self.out, name="decoder")

    def __call__(self, x):
        h = jax.nn.relu(self.encoder(x))
        return self.decoder(h)

For a vector x of length D, choose out=D. The encoder has a (D, hidden) kernel and (hidden,) bias. The decoder has a (hidden, D) kernel and (D,) bias. ReLU sits between the layers; there is no output activation.

Matching the input and output widths does not make the decoder the inverse of the encoder. The two kernels are initialized separately and are not tied as transposes. ReLU can also discard information.

Read the parameter tree

For D=3 and hidden=4, initialization returns this structure:

params
  encoder: kernel (3, 4), bias (4,)
  decoder: kernel (4, 3), bias (3,)

For example, access the encoder kernel with variables["params"]["encoder"]["kernel"]. These paths help identify arrays when inspecting, replacing or loading weights. Matching names alone is not enough to load a checkpoint: shapes and layer conventions must also agree.

Attribute names and explicit names

In this setup-based definition, omitting name="encoder" still gives the unnamed child the name encoder, taken from self.encoder. It does not become Dense_0. An explicit name lets you keep the parameter path even if you later rename the Python attribute. Keep both names as shown here because parameter paths affect seeded initialization.

Setup is useful when you want children available as attributes in multiple forward methods. Compact methods also support explicit names and reuse of a previously created layer instance. Choose the style that makes the calculation readable; naming and sharing are not exclusive to setup.

Let Linen invoke setup when binding the model for initialization or application. The usual unbound model object does not immediately expose setup-created attributes just because its constructor has returned.

Naming is not an optimizer policy

A meaningful path makes it easier to build a parameter mask, but it does not freeze a layer by itself. Optax transformations use the mask you supply; optax.masked does not automatically select parameters by name or freeze all unselected updates. This exercise does no optimization.

Problem

Implement multi_submodule_forward(seed, x, hidden):

  1. Define EncoderDecoder exactly as above, with default Dense settings, including bias.
  2. Set out_size = x.shape[0] and construct EncoderDecoder(hidden=int(hidden), out=int(out_size)).
  3. Initialize with jax.random.PRNGKey(int(seed)) and x, without splitting the key first.
  4. Apply the model with the full initialized variables mapping and x, returning the decoder output directly.

x is a finite float32 JAX vector of length D >= 1. hidden is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; both may arrive as floats. Return a JAX vector of length D, leaving x unchanged. Preserve the child names and do not share their kernels.