Skip to content

← Flax step 71 of 100

Medium Primitives

Repeat Dense Initialization

Jump to the task

Repeating initialization can reproduce a model's initial parameters when the configuration, inputs, random key and relevant software behavior are unchanged. This exercise demonstrates that property for a default Flax Linen Dense layer. It does not load a checkpoint.

Initialize, then initialize again

model = nn.Dense(features)
first = model.init(jax.random.PRNGKey(seed), x)["params"]
second = model.init(jax.random.PRNGKey(seed), x)["params"]
out = model.apply({"params": second}, x)

Both calls use a newly constructed key with the same seed. For this Dense layer in the same environment, the kernel and bias values agree. Passing a different split key can produce a different kernel.

Do not generalize this to arbitrary modules using only seed and input shape. Initialization can depend on module configuration, input values, dtype, RNG streams and library behavior. Reproducibility across software versions or backends is a separate concern.

Reinitialization is not restoration

An optimizer changes parameters after initialization. Calling init again with the original seed recreates initial values; it does not recover those learned updates. Restoring a checkpoint reads the values that were actually saved, potentially along with optimizer state, counters and other variable collections.

Real checkpoint metadata records structure and array information. Restore targets and options can specify the desired types and placement, but a missing template does not universally turn a saved tree into an unnamed flat dictionary. Likewise, saving the full variables tree can be valid: there is no rule that every checkpoint must contain only inner params.

This exercise uses neither Orbax nor filesystem I/O. Matching repeated initializations does not test serialization, restoration, sharding changes or training resumption.

Apply the parameters

For input shape (B, D) and width F, the initialized kernel has shape (D, F) and bias has shape (F,). Applying the layer computes x @ kernel + bias, with shape (B, F). There is no activation function.

Flatten the result in row-major order: all features of the first example, then all features of the second, and so on. Negative outputs remain negative. The final length is B * F.

Problem

Implement orbax_load_simulate(seed, x, features):

  1. Convert seed and features to integers; construct default nn.Dense(F).
  2. Initialize it twice, each time using jax.random.PRNGKey(int(seed)) and the same x. Extract each call's params collection.
  3. Apply the second parameter tree with model.apply({"params": second}, x).
  4. Return the result flattened to a JAX vector of shape (B * F,).

x is a finite float32 JAX matrix of shape (B, D), with positive B and D. Width is a positive integer-valued scalar; seed is a nonnegative integer-valued scalar. Both may arrive as floats. Keep default Dense initializers and bias, and leave x unchanged.

The output tests check the seeded affine calculation. They cannot prove that initialization happened twice, and they do not verify checkpoint loading. Compare both parameter trees locally to check the repeatability step itself.