Skip to content

← Flax step 72 of 100

Hard Primitives

Load Linear Weights into Flax Dense

Jump to the task

Match the layer's matrix convention

To reuse a linear layer's weights, first check how each implementation stores its matrix. For a batch x with shape (N, in_dim), these two layers represent the same affine transformation with different layouts:

Layer Weight shape Computation
PyTorch nn.Linear (out_dim, in_dim) x @ weight.T + bias
Flax Linen nn.Dense (in_dim, out_dim) x @ kernel + bias

Set the Flax kernel to the transpose of the Linear weight. Copy the bias without transposing it: both layers use shape (out_dim,). Reshaping a flat buffer directly to the destination shape is generally different from reshaping to the source shape and then transposing.

The exercise's hf_ argument names refer to supplied weights in this Linear layout. Hugging Face is not a single weight-layout convention; the correct conversion depends on the checkpoint's actual layer types. This task reads no checkpoint files and makes no downloads.

A small conversion you can verify

Let the source weight be [[1, 2, 3], [4, 5, 6]], the bias be [0.5, -0.5], and one input row be [1, 0, -1]. The two output values are 1 - 3 + 0.5 = -1.5 and 4 - 6 - 0.5 = -2.5.

import jax.numpy as jnp
import flax.linen as nn
import numpy as np

source_weight = jnp.array([[1., 2., 3.], [4., 5., 6.]])
bias = jnp.array([0.5, -0.5])
x = jnp.array([[1., 0., -1.]])
variables = {"params": {"kernel": source_weight.T, "bias": bias}}
result = nn.Dense(2).apply(variables, x)
np.testing.assert_allclose(result, [[-1.5, -2.5]])

This supplies every parameter directly, so apply does not need a preceding random initialization. Initializing first can help inspect expected names and shapes, but those random values must not remain in the loaded model.

A square weight matrix needs the same transpose rule. Its shape cannot reveal a missing transpose; compare numerical outputs with the source computation using a nonsymmetric weight matrix and a known input.

Problem

Implement hf_weight_load_dense(seed, x, hf_kernel_flat, hf_bias_flat, features):

  1. Set out_dim = int(features) and read in_dim from x.shape[-1].
  2. Convert the supplied weight values to float32 and reshape them in row-major order to (out_dim, in_dim).
  3. Transpose that matrix to get the Flax (in_dim, out_dim) kernel.
  4. Convert the supplied bias to float32 with shape (out_dim,).
  5. Apply a standalone nn.Dense(out_dim) with these supplied parameters, then return the output flattened in row-major order.

The variables passed to apply must have the structure {"params": {"kernel": converted_kernel, "bias": supplied_bias}}. Do not place an extra Dense_0 level inside this standalone layer.

Inputs:

  • x: float array (N, in_dim), with positive dimensions.
  • hf_kernel_flat: array (out_dim * in_dim,), flattened from the source (out_dim, in_dim) matrix in row-major order.
  • hf_bias_flat: array (out_dim,).
  • features: positive integer-valued output width, possibly a float.
  • seed: integer-valued seed retained for the initialization workflow. The reference initializes with PRNGKey(int(seed)), then replaces all parameters. The final answer therefore does not depend on seed.

Output: a one-dimensional array of length N * out_dim. Both kernel and bias must come from the supplied values. There is no activation, normalization, or training step.

Beyond this one layer

A full model conversion needs rules for its actual parameter names, layer types, packed projections, shapes, and dtypes. Do not transpose every tensor whose name contains weight. A matching array shape alone also cannot establish that two models compute the same function. Compare outputs after conversion with an appropriate numerical tolerance.

Keep seed and features static if JIT-compiling the reference's initialization-based implementation, which converts them to Python integers. See the PyTorch Linear reference for the source matrix convention.