Skip to content

← Flax step 11 of 100

Medium Primitives

Implement Dense from Scratch

Jump to the task

Build a Dense layer using self.param rather than calling nn.Dense. The calculation is an affine map: multiply the input by a learned kernel, then add a learned bias.

Shapes and calculation

For an input vector x of length in_features and features output units:

  • kernel has shape (in_features, features).
  • bias has shape (features,).
  • jnp.dot(x, kernel) + bias has shape (features,).

Each output is a weighted sum of all input features plus its own bias. For example, if x = [2, 3], the kernel is [[1, 4], [2, -1]], and the bias is [0.5, 1], the output is [8.5, 6].

Declare the two parameters

In a compact method, use self.param to create or retrieve an array:

kernel = self.param(
    "kernel", nn.initializers.lecun_normal(),
    (x.shape[-1], self.features),
)
bias = self.param(
    "bias", nn.initializers.zeros, (self.features,),
)

During initialization, Linen supplies an internally derived parameter key to each initializer, followed by the arguments you provided. Here those arguments contain the shape. These initializers have a default dtype, so an explicit dtype argument is not needed for this exercise.

lecun_normal() returns an initializer function. zeros is already an initializer function and should be passed without parentheses. During ordinary apply, self.param retrieves the existing arrays instead of replacing them with new samples.

Since these parameters belong directly to MyDense, the variables are:

{"params": {"kernel": kernel, "bias": bias}}

There is no Dense_0 child entry. Wrapping nn.Dense would introduce a different parameter path and can change the seeded kernel, even though the forward formula is the same.

Initialization and shape checks

JAX's LeCun normal initializer samples a rescaled truncated normal distribution with target variance 1 / in_features. Under independent, zero-mean input and weight assumptions, this helps preserve input variance in a linear layer. It does not guarantee successful training or an exact sample variance. The bias starts at zero.

Read the input width from x.shape[-1]. Initialization executes the forward calculation too, so reversing a nonsquare kernel's dimensions normally causes a matrix-product error during init, not just later during apply. After initialization, applying the variables to a different input feature width also fails because the stored kernel has a fixed shape.

In larger models, choosing a configured input width can be valid. Here, inferring it from x lets the same module definition handle each test's input size without another configuration field.

Problem

Implement MyDense(nn.Module) with a features field and a compact __call__(self, x). Declare kernel and bias exactly as above, then return jnp.dot(x, kernel) + bias with no activation.

Implement my_dense_forward(seed, x, features):

  1. Construct MyDense(features=int(features)).
  2. Initialize it with jax.random.PRNGKey(int(seed)) and x, without splitting the key first.
  3. Apply it with the full initialized variables mapping and x. Return the output directly.

x is a finite float32 JAX vector of length at least one. features is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; either may arrive as a float. Return a JAX vector of length features, leaving x unchanged. Keep the parameter names kernel and bias, and do not create a nested Dense module.