Skip to content

← Flax NNX step 79 of 100

Hard Primitives

NNX: Import PyTorch Linear Weights

Jump to the task

Match the stored weight layout

A PyTorch Linear layer stores its weight as (output_width, input_width) and computes x @ weight.T + bias. An NNX Linear stores its kernel as (input_width, output_width) and computes x @ kernel + bias.

For these two layers, transpose the supplied weight and keep the bias unchanged. This exercise uses an already decoded, flattened array in PyTorch Linear order. It does not download a Hugging Face model, read a checkpoint file or load a complete network.

Restore rows before transposing

Let I be the input width and F the output width. The flat input lists each of the F output rows in order, with I coefficients per row. First reshape to (F, I), then transpose to (I, F).

For example, with I=2 and F=3, the flat values [1, 2, 3, 4, 5, 6] represent these output rows:

PyTorch weight       NNX kernel
[[1, 2],             [[1, 3, 5],
 [3, 4],              [2, 4, 6]]
 [5, 6]]

Reshaping the flat array directly to (I, F) would instead produce [[1, 2, 3], [4, 5, 6]]. Its shape would be valid, but its coefficients would be assigned to the wrong input/output pairs. Shape checks alone cannot catch that error.

A complete transfer example

For x=[2, -1] and bias=[0.5, -1, 2], the three output rows give [2*1 - 2 + 0.5, 2*3 - 4 - 1, 2*5 - 6 + 2], or [0.5, 1, 6].

import jax.numpy as jnp
from flax import nnx

flat = jnp.array([1., 2., 3., 4., 5., 6.])
bias = jnp.array([0.5, -1., 2.])
x = jnp.array([2., -1.])
weight = flat.reshape(3, 2)
model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
model.kernel[...] = weight.T
model.bias[...] = bias
assert jnp.allclose(model(x), jnp.array([0.5, 1., 6.]))
assert jnp.allclose(model(x), weight @ x + bias)

Assign through parameter[...] to keep the nnx.Param wrappers. Both the initialized kernel and bias are replaced, so the destination seed does not affect the final output. Bias remains a vector of length F; there is no bias transpose.

Check the source layer, not just its file name

This transpose rule belongs to PyTorch Linear's storage convention. A model hosted on Hugging Face may contain other layer types or use a different framework. Its weights are not universally (out, in). Embeddings, convolutions, packed projections and quantized tensors need their own layout and representation checks.

Even a square matrix needs the correct orientation. A symmetric weight can hide a missing transpose, whereas an asymmetric example reveals it. A real model port also needs matching parameter names, architecture, configuration and any shared weights or non-parameter state.

Here all arrays are already JAX arrays. The flat representation is an exercise input choice, not a limitation that prevents the site from accepting multidimensional arrays. There is no PyTorch device transfer or checkpoint deserialization to implement.

Problem

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

  1. Set F=int(features), I=x.shape[-1], and construct a biased nnx.Linear(I, F, rngs=nnx.Rngs(int(seed))).
  2. Reshape hf_kernel_flat to (F, I) using its supplied row order, then transpose it to obtain the NNX kernel.
  3. Assign the kernel and hf_bias_flat to the existing parameter wrappers.
  4. Return model(x), shaped (F,).

x is a float32 vector of length I. hf_kernel_flat has F*I float32 values and hf_bias_flat has F float32 values; both are one-dimensional. I and F are positive. seed and features arrive as floats representing integers; keep them static under JIT. Leave the input arrays unchanged.

The returned output checks the affine calculation. It cannot establish by itself that a module was constructed and its weights imported, so follow the requested transfer rather than bypassing it with a dot product.