Skip to content

← Flax step 98 of 100

Medium Primitives

Param Sharing - One Module, Two Call Sites

Jump to the task

Parameter sharing means using the same parameter arrays in more than one part of a calculation. In this exercise, a single Dense layer is applied twice in sequence, with ReLU after each application.

One layer instance, two calls

Define one child module in setup:

def setup(self):
    self.shared = nn.Dense(self.hidden, name="shared")

Then call that same instance twice:

h1 = jax.nn.relu(self.shared(x))
h2 = jax.nn.relu(self.shared(h1))
return h2

Both calls use the same kernel W and bias b. The calculation is relu(relu(x @ W + b) @ W + b). It is not two independently initialized layers and is not generally equivalent to applying one Dense layer once.

The initialized variables contain one child entry:

{"params": {"shared": {"kernel": W, "bias": b}}}

The kernel has shape (hidden, hidden) and the bias has shape (hidden,). That is two array leaves and hidden*hidden + hidden scalar parameters, regardless of the two call sites.

Sharing is about instances, not syntax

setup is convenient for naming a reusable child, but it is not the only way to share parameters. Inside @nn.compact, you can also create a layer once, store it in a local variable and call that variable twice. In contrast, two separate nn.Dense(hidden) constructions create two separate layers.

This exercise asks for the setup version and the explicit name "shared". Keep that name because parameter paths affect random initialization. Let Flax invoke setup; do not call it manually.

Why the feature widths must match

The first call sees input width hidden and creates a square kernel. Its output also has width hidden, so the second call can reuse that kernel. If the original input had a different width, the second call would require a different kernel shape and fail.

This square-kernel requirement comes from applying the layer sequentially to its own output. Sharing a Dense across two separate inputs only requires that both inputs have the same feature width; a shared kernel need not always be square.

Problem

Implement param_sharing_forward(seed, x, hidden):

  1. Define SharedDense(nn.Module) with a hidden field and the single named Dense child shown above. Use default Dense settings and bias.
  2. Apply that child twice with ReLU after each call.
  3. Initialize SharedDense(hidden=int(hidden)) using jax.random.PRNGKey(int(seed)) and x, without splitting the key first.
  4. Apply the model with the whole initialized variables mapping and x. Return the output flattened in row-major order with .reshape(-1).

x is a finite float32 JAX array of shape (N, hidden), with N >= 1. hidden is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; either may arrive as a float. Return a JAX array of shape (N * hidden,), leaving x unchanged. Do not train or update the parameters. You can inspect the params mapping to confirm that both calls use its one shared kernel and bias.