Skip to content

← Flax step 70 of 100

Medium Primitives

Count Parameter Arrays in a Dense Stack

Jump to the task

A parameter tree contains arrays organized by layer and parameter name. Counting its leaves tells you how many arrays it contains. It does not tell you how many scalar parameters, bytes or checkpoint files there are. This exercise practices inspecting a Flax Linen parameter tree.

Build separate layers

In a compact module, create a new Dense instance on each iteration:

class StackedDense(nn.Module):
    features: int
    num_layers: int

    @nn.compact
    def __call__(self, x):
        for _ in range(self.num_layers):
            x = nn.Dense(self.features)(x)
        return x

Each layer has its own kernel and bias. With width D, their shapes are (D, D) and (D,). There are no activation functions between layers. The automatic parameter names are Dense_0, Dense_1, and so on.

Creating one Dense instance and reusing it is different: its calls share parameters. Likewise, repeating a reference with [layer] * N does not construct N independent layers.

Inspect the parameter tree

variables = model.init(jax.random.PRNGKey(seed), x)
params = variables["params"]
leaves = jax.tree_util.tree_leaves(params)
leaf_count = len(leaves)

A kernel array is one leaf, however many elements it contains. A bias array is another. This model therefore has 2 * num_layers leaves. For three width-four layers, that is 6 array leaves but 3 * (4 * 4 + 4) = 60 scalar parameters.

Leaf counting is one useful inspection step. Equal counts do not prove that two trees have matching names, shapes, dtypes or values.

Relation to checkpointing

Flax Linen initialization returns variable collections. This model uses only params; other models can also have collections such as running statistics. A complete training checkpoint may additionally need optimizer, step and RNG state. Saving only parameters is not always enough to resume training.

Orbax can serialize state trees, but storage layout depends on the format and options. One leaf does not imply one file. This task makes no Orbax calls and writes no files. It checks tree inspection, not serialization, checkpoint completeness or restore correctness.

Problem

Implement orbax_save_and_count(seed, x, num_layers):

  1. Convert seed and layer count to integers. Set D = x.shape[-1].
  2. Construct the compact stack above with num_layers separate default Dense layers, each of width D and with bias enabled.
  3. Initialize using jax.random.PRNGKey(int(seed)) and x.
  4. Count jax.tree_util.tree_leaves(variables["params"]).
  5. Return a floating-point JAX array of shape (1,): [leaf_count].

x is a finite float32 JAX matrix of shape (B, D), with positive B and D. Layer count is a positive integer-valued scalar and seed is nonnegative and integer-valued; either may arrive as a float. Leave x unchanged. Neither batch size, width, input values nor seed changes the leaf-count formula for this model.

The output tests verify the count. They cannot establish that a submission built the model rather than returning 2 * num_layers.