LazyLinear: infer in_features from the first real forward call
The problem, from first principles
[01-linear-forward]'s linear_forward needs weight to already exist, with a specific in_features baked into its shape, before it can run at all. That's a real inconvenience when building a network: in_features for the FIRST layer of a network usually depends on the shape of your actual dataset (how many pixels an image has after flattening, how many features a tabular dataset has after preprocessing), which you may not want to compute and pass in by hand every time you experiment with a new dataset or a different preprocessing pipeline. LazyLinear removes this friction: you only ever specify out_features up front, and in_features gets figured out automatically, the FIRST time real data actually flows through the layer.
The tricky part isn't the inference itself (x.shape[-1] tells you in_features immediately), it's making sure the layer only does this ONCE: every forward call after the first one must reuse the exact same weight and bias it created on that first call, not silently reinitialize (and therefore erase any training progress) on every subsequent call.
From theory to code
Implement LazyLinear.forward(self, x). self.weight and self.bias start as None (already set up in __init__, along with self.weight_init/self.bias_init, the functions to call for generating fresh initial values). On the FIRST call to forward, if self.weight is None, infer in_features = x.shape[-1], call self.weight_init(in_features, self.out_features) and self.bias_init(self.out_features) to create the actual arrays, register them as parameters (via the inherited self.register_parameter, from [05-module-base-class]), and store them on self. On every call (first and all subsequent), delegate the actual computation to linear_forward (already provided, reused from [01-linear-forward]).
Constraints
- The very first call to
forwardmust createself.weightwith shape(out_features, in_features)andself.biaswith shape(out_features,), usingx's LAST dimension asin_features. - Every subsequent call must reuse the SAME
weight/biasobjects created on the first call, never re-running the initializers again. - The actual output computation, on every call, must delegate to
linear_forward, not reimplement it.
Hints
Hint 1: The "only once" check
if self.weight is None: is exactly the right guard: it's None only before the very first call, and becomes a real array immediately after, so this check alone prevents any re-initialization on later calls.
Hint 2: Inferring in_features
x.shape[-1] is the size of x's LAST dimension, which for a batch of shape (batch_size, in_features) is exactly in_features, regardless of what batch_size happens to be.
Hint 3: Don't forget to register
After creating self.weight and self.bias inside the if block, call self.register_parameter("weight", self.weight) and self.register_parameter("bias", self.bias) so .parameters() (inherited from Module) picks them up, exactly like a normal, non-lazy layer would.
Click "Run Tests" to test your implementation