from dataclasses import dataclass
from d2l import torch as d2l
import torch
from torch import nn
from torch.nn import functional as F5.1 Modules and Model Construction
The networks we have trained so far were small enough to write down layer by layer. The networks ahead are not: deep ResNets chain more than a hundred convolutional layers (He et al. 2016), and a GPT-style language model stacks dozens of identical Transformer blocks (Radford et al. 2019). Such models are assembled from repeated blocks, groups of layers treated as units of design. The abstraction that makes blocks composable is the module.
A module is an object with three responsibilities: it owns parameters, it owns child modules, and it implements a forward computation that maps inputs to outputs. The definition is deliberately recursive. A fully connected layer is a module (parameters, no children). A residual block is a module (no direct parameters, a few child layers that own parameters). A hundred-layer network is a module whose children are blocks whose children are layers. Most models therefore have a tree-shaped module hierarchy, as sketched in Figure 5.1.1. Reusing one child at several sites turns that tree into an object graph, a case we meet when tying parameters in Section 5.2. Almost everything this chapter does to a model, listing its parameters (Section 5.2), moving it to a GPU (Section 5.7), saving it to disk (Section 5.6), is implemented as a walk over that tree.
from dataclasses import dataclass
import tensorflow as tffrom dataclasses import dataclass
from d2l import jax as d2l
import jax
from jax import numpy as jnp
from flax import nnxfrom dataclasses import dataclass
from mxnet import init, np, npx
from mxnet.gluon import nn
npx.set_np()5.1.1 The Module Abstraction
In PyTorch the module class is nn.Module. We have used one of its subclasses all along: nn.Sequential builds a model from a chain of layers, here the familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output layer.
In TensorFlow the module class is Keras’s tf.keras.Model. We have used one of its subclasses all along: tf.keras.Sequential builds a model from a list of layers, here the familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output layer. Keras attaches the activation to the Dense layer itself rather than interposing a separate activation layer.
In JAX the module class is Flax’s nnx.Module. We have used one of its subclasses all along: nnx.Sequential builds a model from a chain of layers, here the familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output layer. NNX modules own their initialized parameters, so construction requires the layer widths and an RNG stream; a forward pass then uses the same direct net(X) notation as the other frameworks.
In MXNet the module class is Gluon’s nn.Block. We have used one of its subclasses all along: nn.Sequential builds a model from a chain of layers appended with its add method, here the familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output layer. Gluon attaches the activation to the Dense layer itself rather than interposing a separate activation layer, and one initialize() call stands between construction and the first forward pass; what that call does, and what it defers, is where this section ends.
net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))
X = torch.rand(2, 20)
net(X).shapetorch.Size([2, 10])net = tf.keras.Sequential([tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(10)])
X = tf.random.uniform((2, 20))
net(X).shapeTensorShape([2, 10])net = nnx.Sequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,
nnx.Linear(256, 10, rngs=nnx.Rngs(1)))
X = jax.random.uniform(d2l.get_key(), (2, 20))
net(X).shape(2, 10)net = nn.Sequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))
net.initialize()
X = np.random.uniform(size=(2, 20))
net(X).shape(2, 10)Sequential is not a special construct. It is itself a module whose forward computation runs its children in order, and the children are the three modules we passed to it, stored under names in a registry:
Sequential is not a special construct. It is itself a model whose forward computation runs its children in order, and the children are the two layers we passed to it, held in a tracked list:
Sequential is not a special construct. It is itself a module whose forward computation runs its children in order, and the children are the three modules we passed to it. Each child owns its parameters; selecting nnx.Param state gives a flat path view whose structure mirrors the module tree:
Sequential is not a special construct. It is itself a Block whose forward computation runs its children in order, and the children are the two layers add appended, stored under names in a registry that printing the model displays:
net._modules{'0': Linear(in_features=20, out_features=256, bias=True),
'1': ReLU(),
'2': Linear(in_features=256, out_features=10, bias=True)}net.layers[<Dense name=dense, built=True>, <Dense name=dense_1, built=True>][(path, value.shape)
for path, value in nnx.state(net, nnx.Param).flat_state()][(('layers', 0, 'bias'), (256,)),
(('layers', 0, 'kernel'), (20, 256)),
(('layers', 2, 'bias'), (10,)),
(('layers', 2, 'kernel'), (256, 10))]netSequential(
(0): Dense(20 -> 256, Activation(relu))
(1): Dense(256 -> 10, linear)
)This registry is what the nn.Module machinery traverses: net.parameters() collects parameters by walking _modules recursively, and the same walk underlies device movement and serialization. A module the registry does not contain might as well not exist, a fact we will exploit for a demonstration shortly.
This list is what the Keras machinery traverses: net.trainable_variables collects variables by walking the children recursively, and the same walk underlies serialization. A layer the tracker does not see might as well not exist, though as we will see shortly, Keras goes to some lengths to make sure that cannot happen by accident.
nnx.state(net, nnx.Param) selects the trainable leaves of the object graph. Optimizer updates, device movement, and serialization traverse this state. The parameters remain owned by their layers; filters give us a state view without maintaining a second model representation by hand.
This registry is what the Block machinery traverses: net.collect_params() gathers parameters by walking the children recursively and keys them by those registry names (0.weight, 1.bias, …), and the same walk underlies initialization, device movement, and serialization. A block the registry does not contain might as well not exist, a fact we will exploit for a demonstration shortly.
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.hidden = nn.LazyLinear(256)
self.out = nn.LazyLinear(10)
def forward(self, X):
return self.out(F.relu(self.hidden(X)))class MLP(tf.keras.Model):
def __init__(self):
super().__init__()
self.hidden = tf.keras.layers.Dense(256, activation='relu')
self.out = tf.keras.layers.Dense(10)
def call(self, X):
return self.out(self.hidden(X))class MLP(nnx.Module):
def __init__(self, rngs):
self.hidden = nnx.Linear(20, 256, rngs=rngs)
self.out = nnx.Linear(256, 10, rngs=rngs)
def __call__(self, X):
return self.out(nnx.relu(self.hidden(X)))class MLP(nn.Block):
def __init__(self):
super().__init__()
self.hidden = nn.Dense(256, activation='relu')
self.out = nn.Dense(10)
def forward(self, X):
return self.out(self.hidden(X))net = MLP()
net(X).shapetorch.Size([2, 10])net = MLP()
net(X).shapeTensorShape([2, 10])net = MLP(nnx.Rngs(0))
net(X).shape(2, 10)net = MLP()
net.initialize()
net(X).shape(2, 10)Two details do the work here. First, self.hidden = nn.LazyLinear(256) is not an ordinary attribute assignment: nn.Module intercepts __setattr__, sees that the value is a module, and adds it to the registry we just inspected. That is why both layers’ parameters show up in net.parameters() with no further ceremony. Second, we never wrote a backward method; automatic differentiation derives gradients from whatever forward computes.
Two details do the work here. First, self.hidden = tf.keras.layers.Dense(...) is not an ordinary attribute assignment: Keras intercepts __setattr__, sees that the value is a layer, and adds it to the tracked children we just inspected. That is why both layers’ variables show up in net.trainable_variables with no further ceremony. Second, we never wrote a backward method; automatic differentiation derives gradients from whatever call computes.
Two details do the work here. First, assigning nnx.Linear children to self.hidden and self.out makes them part of the NNX object graph, and their attribute names become paths in the selected parameter state. Second, we never wrote a backward method; nnx.grad derives gradients from whatever __call__ computes. Layers are constructed once in __init__, keeping parameter creation separate from the forward computation.
Two details do the work here. First, self.hidden = nn.Dense(...) is not an ordinary attribute assignment: nn.Block intercepts __setattr__, sees that the value is a block, and adds it to the registry we just inspected. That is why both layers’ parameters show up in net.collect_params() with no further ceremony. Second, we never wrote a backward method; automatic differentiation derives gradients from whatever forward computes.
5.1.2 Sequential and Friends: Containers
To see that there is no magic left in nnx.Sequential, we can write it ourselves. Two ingredients suffice: register each child under a name, and loop over the children in forward.
To see that there is no magic left in Sequential, we can write it ourselves. Two ingredients suffice: store the children in an attribute, and loop over them in call.
To see that there is no magic left in nnx.Sequential, we can write it ourselves. Two ingredients suffice: declare a field that holds the list of children, and loop over them in __call__.
To see that there is no magic left in nnx.Sequential, we can write it ourselves. Two ingredients suffice: register each child, and loop over the children in forward. One Gluon 2.0 wrinkle: the registry holds weak references, so our class also keeps the blocks in a plain list to keep them alive.
class MySequential(nn.Module):
def __init__(self, *args):
super().__init__()
for idx, module in enumerate(args):
self.add_module(str(idx), module)
def forward(self, X):
for module in self.children():
X = module(X)
return Xclass MySequential(tf.keras.Model):
def __init__(self, *args):
super().__init__()
self.modules = args
def call(self, X):
for module in self.modules:
X = module(X)
return Xclass MySequential(nnx.Module):
def __init__(self, *modules):
self.modules = nnx.List(modules)
def __call__(self, X):
for module in self.modules:
X = module(X)
return Xclass MySequential(nn.Block):
def __init__(self):
super().__init__()
# The registry holds weakrefs; this list keeps the blocks alive
self._layers = []
def add(self, block):
self._layers.append(block)
self.register_child(block)
def forward(self, X):
for block in self._layers:
X = block(X)
return Xadd_module writes a child into the registry under a string name (that is where the '0', '1', '2' keys above came from), and children() iterates the registry in insertion order. Our version is a drop-in replacement:
self.modules = args looks like an ordinary assignment, and that is the point: the __setattr__ interception scans whatever is assigned, looking inside lists, tuples, and dictionaries, so both Dense children are tracked and appear in net.layers. Our version is a drop-in replacement:
nnx.List is a graph-aware container. Both Linear children are tracked (nnx.relu is a plain function, with nothing to track), and their parameters appear beneath the modules path. Our version is a drop-in replacement:
register_child writes a child into the registry under a string name (that is where the (0) and (1) keys above came from; attribute assignment calls it with the attribute name instead). Our version is a drop-in replacement:
net = MySequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))
net(X).shapetorch.Size([2, 10])net = MySequential(tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(10))
net(X).shapeTensorShape([2, 10])net = MySequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,
nnx.Linear(256, 10, rngs=nnx.Rngs(1)))
net(X).shape(2, 10)net = MySequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))
net.initialize()
net(X).shape(2, 10)The registration step is easy to lose. The following module looks reasonable, and its forward pass works, so nothing appears wrong:
In older imperative frameworks this registration step is famously easy to lose: store the children in a plain Python list instead of the framework’s dedicated container, and their parameters silently vanish from the model. Keras closes that trap. Because the attribute scan looks inside lists and dictionaries, a plain list of layers is tracked like any other child:
Graph-aware containers make registration explicit in NNX. nnx.List tracks the child modules and still supports ordinary indexing and iteration:
The registration step is easy to lose. The following block stores its layers in a plain Python list, which the __setattr__ interception ignores, so nothing inside it is registered. Since net.initialize() walks the registry, it cannot reach the hidden children either, so we initialize each layer by hand just to get a forward pass, and the model appears to work:
class PlainListMLP(nn.Module):
def __init__(self):
super().__init__()
self.layers = [nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10)]
def forward(self, X):
for layer in self.layers:
X = layer(X)
return X
net = PlainListMLP()
net(X).shape, sum(p.numel() for p in net.parameters())(torch.Size([2, 10]), 0)class ListMLP(tf.keras.Model):
def __init__(self):
super().__init__()
self.blocks = [tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(10)]
def call(self, X):
for block in self.blocks:
X = block(X)
return X
net = ListMLP()
net(X).shape, sum(int(tf.size(v)) for v in net.trainable_variables)(TensorShape([2, 10]), 7946)class ListMLP(nnx.Module):
def __init__(self, rngs):
self.layers = nnx.List([nnx.Linear(20, 256, rngs=rngs),
nnx.Linear(256, 10, rngs=rngs)])
def __call__(self, X):
return self.layers[1](nnx.relu(self.layers[0](X)))
net = ListMLP(nnx.Rngs(0))
state = nnx.state(net, nnx.Param)
net(X).shape, sum(x.size for _, x in state.flat_state())((2, 10), 7946)class PlainListMLP(nn.Block):
def __init__(self):
super().__init__()
self.layers = [nn.Dense(256, activation='relu'), nn.Dense(10)]
for layer in self.layers:
layer.initialize() # By hand: net.initialize() cannot see them
def forward(self, X):
for layer in self.layers:
X = layer(X)
return X
net = PlainListMLP()
net(X).shape, len(net.collect_params())/tmp/ipykernel_65430/1272845916.py:14: UserWarning: "PlainListMLP.layers" is an unregistered container with Blocks. Note that Blocks inside the list, tuple or dict will not be registered automatically. Make sure to register them using register_child() or switching to nn.Sequential/nn.HybridSequential instead.
net(X).shape, len(net.collect_params())((2, 10), 0)The model computes, yet it owns zero parameters. A plain Python list is not a module, so the __setattr__ interception ignores it and nothing inside it is registered. The layers still hold perfectly good tensors, which is why forward runs. But net.parameters() is empty: hand this model to an optimizer and training proceeds without a single error while updating nothing; net.to(device) leaves every layer behind on the CPU; net.state_dict() checkpoints an empty model. Because no step raises an exception, this bug is usually diagnosed by staring at a loss curve that refuses to move.
All 7946 parameters are present; there is no broken variant of this model to show. (We named the list blocks only because layers is reserved: Keras maintains that attribute itself, and assigning to it raises an error.) The structural mistake Keras does guard against is different in kind and, more usefully, in loudness. Building is a commitment: once a model’s variables exist, its set of children is locked, and attaching a new layer to a built model fails at the assignment itself:
All 7946 parameters are present. NNX deliberately rejects a plain Python list that contains graph nodes, because treating such a container as static would silently hide its parameters. The error points directly to nnx.List (or nnx.data) as the repair:
The model computes, yet it owns zero parameters: hand it to a Trainer and training updates nothing; save_parameters checkpoints an empty model. Gluon does not let this pass in silence, though. collect_params() scans the block’s ordinary attributes for lists, tuples, and dictionaries holding unregistered blocks, and the call above printed a warning naming the guilty attribute ("PlainListMLP.layers" is an unregistered container with Blocks...) along with the fix. And had we skipped the by-hand initialization, the first forward pass would have stopped with a Parameter has not been initialized error instead of running at all. Where PyTorch’s version of this bug is diagnosed by staring at a loss curve, Gluon’s announces itself.
class ModuleListMLP(PlainListMLP):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList(self.layers)
net = ModuleListMLP()
net(X).shape, sum(p.numel() for p in net.parameters())(torch.Size([2, 10]), 7946)try:
net.head = tf.keras.layers.Dense(2) # net is built: too late
except ValueError as e:
print(e)You cannot add new elements of state (variables or sub-layers) to a layer that is already built. All state must be created in the `__init__()` method or in the `build()` method.class PlainListSequential(nnx.Module):
def __init__(self, *modules):
self.modules = list(modules) # Graph nodes cannot be static data
try:
net = PlainListSequential(nnx.Linear(20, 10, rngs=nnx.Rngs(0)))
except ValueError as e:
print(e)Found data on value of type '<class 'list'>' assigned to static attribute 'modules' of Pytree type '<class '__main__.PlainListSequential'>'. Static attributes should not contain data values. Consider one of the following options:
1. If the attribute is meant to be static, remove the data values of type(s):
Linear
2. If the attribute is meant to be data, wrap the value with nnx.data on assignment:
_.modules = nnx.data(...)
3. Alternatively, annotate the class attribute with nnx.data:
class PlainListSequential(Module):
modules: list = nnx.data()
4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.
5. Disable pytree for this class:
class PlainListSequential(Module, pytree=False):class RegisteredListMLP(PlainListMLP):
def __init__(self):
super().__init__()
for layer in self.layers:
self.register_child(layer)
net = RegisteredListMLP()
net(X).shape, sum(p.data().size for p in net.collect_params().values())((2, 10), 7946)Same forward pass, 7946 registered parameters. The division of labor among the containers is now clear. nn.Sequential registers its children and supplies the run-them-in-order forward. nn.ModuleList registers a list of children but supplies no forward at all; you keep writing the loop, which is exactly what you want when the loop body is not a plain chain (blocks that take extra arguments, or a skip connection around each block). nn.ModuleDict does the same for children indexed by name. Transformer implementations conventionally keep their stack of blocks in an nn.ModuleList and their named parts (embedding, final normalization, output head) as attributes.
The error message states the rule Keras enforces: all state must be created in __init__ or in build, never after. Keras therefore ships no special list or dict containers, because none are needed: any layer reachable from an attribute is tracked, and a late structural edit fails at assignment time rather than being silently ignored. Transformer implementations in Keras conventionally keep their stack of blocks in exactly the kind of plain list ListMLP used, with the named parts (embedding, final normalization, output head) as separate attributes.
NNX supplies nnx.List and nnx.Dict for graph nodes. A misdeclared child fails at construction time instead of yielding a model that runs but trains nothing. Transformer implementations conventionally keep their stack of blocks in an nnx.List, with the embedding and output head as named attributes.
Same forward pass, 7946 registered parameters, and no warning. Gluon ships no ModuleList or ModuleDict equivalents: the two registered homes for children are named attributes and nn.Sequential, with register_child as the escape hatch for anything else, such as a list of blocks whose loop body is not a plain chain. Transformer implementations in Gluon conventionally keep their stack of blocks in an nn.Sequential (or register them one by one) and their named parts (embedding, final normalization, output head) as attributes.
5.1.3 Forward Is Just Python
forward is an ordinary Python method. Nothing restricts it to chaining children: it can branch, loop, call any tensor function, and combine intermediate results however it likes. The loop in ModuleListMLP already used this freedom. Its most consequential one-line use is the residual connection, the wiring idiom at the heart of ResNets and Transformers alike:
call is an ordinary Python method. Nothing restricts it to chaining children: it can branch, loop, call any TensorFlow function, and combine intermediate results however it likes; TensorFlow executes eagerly, so all of this runs one operation at a time, just as in NumPy. (Once a model is wrapped in tf.function for speed, as the Trainer from Section 2.2 does, AutoGraph rewrites such control flow into graph form.) The loop in MySequential already used this freedom. Its most consequential one-line use is the residual connection, the wiring idiom at the heart of ResNets and Transformers alike:
__call__ is an ordinary Python method. Nothing restricts it to chaining children: it can branch, loop, call any jnp function, and combine intermediate results however it likes. The loop in MySequential already used this freedom. Its most consequential one-line use is the residual connection, the wiring idiom at the heart of ResNets and Transformers alike. The residual body is constructed once and then reused on every call.
forward is an ordinary Python method. Nothing restricts it to chaining children: it can branch, loop, call any np function, and combine intermediate results however it likes; a Gluon Block executes eagerly, so all of this runs one operation at a time, just as in NumPy. (Gluon’s HybridBlock can compile the forward computation into a graph for speed, at the price of restricting it to traceable operations; we stay with Block here.) The loop in MySequential already used this freedom. Its most consequential one-line use is the residual connection, the wiring idiom at the heart of ResNets and Transformers alike:
class ResidualBlock(nn.Module):
def __init__(self, num_hiddens):
super().__init__()
self.body = nn.Sequential(
nn.Linear(num_hiddens, num_hiddens), nn.ReLU(),
nn.Linear(num_hiddens, num_hiddens))
def forward(self, X):
return X + self.body(X)class ResidualBlock(tf.keras.Model):
def __init__(self, num_hiddens):
super().__init__()
self.body = tf.keras.Sequential([
tf.keras.layers.Dense(num_hiddens, activation='relu'),
tf.keras.layers.Dense(num_hiddens)])
def call(self, X):
return X + self.body(X)class ResidualBlock(nnx.Module):
def __init__(self, num_hiddens, rngs):
self.body = nnx.Sequential(
nnx.Linear(num_hiddens, num_hiddens, rngs=rngs), nnx.relu,
nnx.Linear(num_hiddens, num_hiddens, rngs=rngs))
def __call__(self, X):
return X + self.body(X)class ResidualBlock(nn.Block):
def __init__(self, num_hiddens):
super().__init__()
self.body = nn.Sequential()
self.body.add(nn.Dense(num_hiddens, activation='relu'),
nn.Dense(num_hiddens))
def forward(self, X):
return X + self.body(X)
X + body(X): the input splits at a branch point into the body stack and an identity skip, and the two rejoin by addition before the block’s output.
X + body(X) is not a layer any framework provides. It is plain arithmetic in the forward computation, and it changes what the block is: the block computes a perturbation of the identity function rather than an arbitrary transformation. If its body is \(F\), its Jacobian is \(I + J_F\), so backpropagation receives an additive identity contribution along the skip path. This contribution can still cancel against \(J_F\); it is a direct path, not a guarantee that every gradient is preserved (He et al. 2016). Figure 5.1.2 diagrams exactly this wiring, and Section 7.4 develops both points when we build ResNet; for now we only need the mechanics. One mechanical consequence is visible already: the addition forces the input and output shapes to agree, so a residual block has a single width that is part of its identity.
That is why we gave body explicit nn.Linear layers rather than lazy ones.
Keras always infers input widths at build time, so here the constraint falls on the caller: feed the block anything other than num_hiddens columns and the addition fails.
That is why num_hiddens is an explicit constructor argument rather than a width inferred from a sample batch.
Gluon defers input widths to the first forward pass (a mechanism we examine at the end of this section), so here the constraint falls on the caller: feed the block anything other than num_hiddens columns and the addition fails.
block = ResidualBlock(24)
block(torch.randn(2, 24)).shapetorch.Size([2, 24])block = ResidualBlock(24)
block(tf.random.normal((2, 24))).shapeTensorShape([2, 24])block = ResidualBlock(24, nnx.Rngs(0))
X24 = jax.random.normal(d2l.get_key(), (2, 24))
block(X24).shape(2, 24)block = ResidualBlock(24)
block.initialize()
block(np.random.normal(size=(2, 24))).shape(2, 24)forward may also use state that is neither an input nor a parameter. Suppose we want to damp each block’s contribution by a fixed factor:
call may also use state that is neither an input nor a parameter. Suppose we want to damp each block’s contribution by a fixed factor:
__call__ may also use state that is neither an input nor a parameter. Suppose we want to damp each block’s contribution by a fixed factor:
forward may also use state that is neither an input nor a parameter. Suppose we want to damp each block’s contribution by a fixed factor:
class ScaledResidual(ResidualBlock):
def __init__(self, num_hiddens, alpha=0.5):
super().__init__(num_hiddens)
self.alpha = torch.tensor(alpha) # Fixed by design, never trained
def forward(self, X):
return X + self.alpha * self.body(X)
block = ScaledResidual(24)
'alpha' in block.state_dict(), list(block.state_dict())[:2](False, ['body.0.weight', 'body.0.bias'])class ScaledResidual(ResidualBlock):
def __init__(self, num_hiddens, alpha=0.5):
super().__init__(num_hiddens)
self.alpha = tf.constant(alpha) # Fixed by design, never trained
def call(self, X):
return X + self.alpha * self.body(X)
block = ScaledResidual(24)
block(tf.random.normal((2, 24)))
any('alpha' in w.path for w in block.weights), [w.path for w in block.weights][:2](False,
['scaled_residual/sequential_2/dense_11/kernel',
'scaled_residual/sequential_2/dense_11/bias'])class ScaledResidual(ResidualBlock):
def __init__(self, num_hiddens, rngs, alpha=0.5):
super().__init__(num_hiddens, rngs)
self.alpha = alpha # Fixed architecture data, never trained
def __call__(self, X):
return X + self.alpha * self.body(X)
block = ScaledResidual(24, nnx.Rngs(0))
block.alpha, [path for path, _ in nnx.state(block, nnx.Param).flat_state()](0.5,
[('body', 'layers', 0, 'bias'),
('body', 'layers', 0, 'kernel'),
('body', 'layers', 2, 'bias'),
('body', 'layers', 2, 'kernel')])class ScaledResidual(ResidualBlock):
def __init__(self, num_hiddens, alpha=0.5):
super().__init__(num_hiddens)
self.alpha = np.array(alpha) # Fixed by design, never trained
def forward(self, X):
return X + self.alpha * self.body(X)
block = ScaledResidual(24)
'alpha' in block.collect_params(), list(block.collect_params())[:2](False, ['body.0.weight', 'body.0.bias'])alpha enters the computation, but it is not a parameter: it never appears in named_parameters(), so the optimizer never touches it. That much we wanted. Storing it as a plain attribute has a cost we did not want, though: as the output shows, it is missing from state_dict() as well, so it will not be saved with the model, and .to(device) will not move it. Some state is not a parameter but must still travel with the model; the registered home for such state is a buffer, introduced in Section 5.2.
alpha enters the computation, but it is not a parameter: as the output shows, block.weights contains only the Dense variables, so the optimizer never touches it. That much we wanted. Storing it as a tf.constant has a cost we did not want, though: not being a variable, it is invisible to weight checkpointing as well, so it will not be saved with the model. Some state is not a parameter but must still travel with the model; the registered home for such state is a non-trainable tf.Variable, introduced in Section 5.2.
alpha enters the computation, but it is ordinary architecture data rather than an nnx.Param, so the optimizer never touches it. Constructing ScaledResidual(24, nnx.Rngs(0), alpha=0.5) reproduces it. Mutable non-parameter state, such as running statistics, instead uses another nnx.Variable subclass, introduced in Section 5.2.
alpha enters the computation, but it is not a parameter: as the output shows, it never appears in collect_params(), so the Trainer never touches it. That much we wanted. Storing it as a plain attribute has a cost we did not want, though: invisible to the parameter walk, it will not be saved by save_parameters, and it will not move to a GPU with the rest of the block. Some state is not a parameter but must still travel with the model; the registered home for such state is a constant parameter (gluon.Constant), introduced in Section 5.2.
class HalvingMLP(nnx.Module):
def __init__(self, rngs):
self.proj = nnx.Linear(20, 24, rngs=rngs)
def __call__(self, X):
X = self.proj(X)
while jnp.abs(X).sum() > 1: # Ordinary Python control flow
X = X / 2
return X.sum()
net = HalvingMLP(nnx.Rngs(0))
net(X)Array(-0.05208473, dtype=float32)Once a model is wrapped in jax.jit for speed, such data-dependent Python control flow must be expressed with jax.lax primitives instead; that constraint arrives only with the compiler, not with the module abstraction.
5.1.4 Lazy Initialization: Shapes from Data
We have been doing something odd since our first MLP without commenting on it: nn.LazyLinear(256) names only the layer’s output width. Its weight matrix has shape (256, in_features), and we never said what in_features is. The layer cannot know it at construction time, since it depends on the data it will receive. So it does not allocate parameters at construction time at all:
We have been doing something odd since our first MLP without commenting on it: Dense(256) names only the layer’s output width. Its kernel has shape (in_features, 256), and we never said what in_features is. The layer cannot know it at construction time, since it depends on the data it will receive. So Keras never allocates variables at construction time: every layer has a build(input_shape) method that creates them, and __call__ invokes it the first time data arrives. Deferred building is not a special mode; it is how every Keras layer works.
NNX makes a different choice: nnx.Linear takes both its input and output widths and creates its parameters in the constructor. This keeps the model object complete: an optimizer or checkpoint can inspect every parameter immediately, without a dummy forward pass. The widths also make a bad connection easier to diagnose, although an ordinary builder does not validate adjacent layers; an incompatible pair still fails when the model is called.
We have been doing something odd since our first MLP without commenting on it: nn.Dense(256) names only the layer’s output width. Its weight matrix has shape (256, in_units), and we never said what in_units is. The layer cannot know it at construction time, since it depends on the data it will receive. So Gluon does not allocate parameters at construction time at all. Deferred initialization is not a special mode or a dedicated Lazy* class; it is how every Gluon layer works, and in_units is simply an optional argument:
net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))
net[0].weight<UninitializedParameter>net = tf.keras.Sequential([tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(10)])
net.weights[]net = nn.Sequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))
net.collect_params(){'0.weight': Parameter (shape=(256, -1), dtype=float32),
'0.bias': Parameter (shape=(256,), dtype=float32),
'1.weight': Parameter (shape=(10, -1), dtype=float32),
'1.bias': Parameter (shape=(10,), dtype=float32)}The weight is a placeholder. The first time data flows through, the layer reads the input width from the batch, allocates and initializes a real weight, and replaces itself with a plain nn.Linear. Its output width then fixes the input of the next lazy layer, and shapes cascade through the whole model:
There are no variables at all yet, and accessing net.layers[0].kernel would raise an error telling us to build the layer first. The first time data flows through, each layer’s build reads the input width from the incoming shape and allocates and initializes a real kernel, and the layer’s declared output width fixes the input of the next layer, so shapes cascade through the whole model:
There is no uninitialized form to inspect. We state the two boundary widths and the hidden width once; the resulting state is available immediately:
Each weight is a placeholder: the parameter objects exist, but the special value -1 marks the input dimension as unknown, and reading net[0].weight.data() would raise an error telling us to initialize the network first. Even initialize() allocates nothing at this point; it only records which initializer (and device) to use once the shapes are known. The first time data flows through, each layer reads its input width from the incoming shape, allocates and initializes a real weight, and its declared output width fixes the input of the next layer, so shapes cascade through the whole model:
net(X)
net[0].weight.shapetorch.Size([256, 20])net(X)
[w.shape for w in net.weights][TensorShape([20, 256]),
TensorShape([256]),
TensorShape([256, 10]),
TensorShape([10])]net = nnx.Sequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,
nnx.Linear(256, 10, rngs=nnx.Rngs(1)))
[(path, value.shape)
for path, value in nnx.state(net, nnx.Param).flat_state()][(('layers', 0, 'bias'), (256,)),
(('layers', 0, 'kernel'), (20, 256)),
(('layers', 2, 'bias'), (10,)),
(('layers', 2, 'kernel'), (256, 10))]net.initialize()
net(X)
net.collect_params(){'0.weight': Parameter (shape=(256, 20), dtype=float32),
'0.bias': Parameter (shape=(256,), dtype=float32),
'1.weight': Parameter (shape=(10, 256), dtype=float32),
'1.bias': Parameter (shape=(10,), dtype=float32)}Lazy layers can remove shape arithmetic from model definitions. In a convolutional network the flattened feature-map size depends on the input resolution and upstream strides. We usually avoid that fragile arithmetic by using global or adaptive pooling before the first linear layer. When a width is part of the architecture, we name it in the model config and pass it to adjacent layers — code that stays explicit and compact without hiding parameter creation behind a sample batch, whether or not the framework offers a lazy mode.
Before the first forward pass, lazy layers contain registered UninitializedParameter placeholders rather than shaped arrays. An optimizer can be constructed over those placeholders, but reading shapes, counting elements, or applying a shape-dependent initializer requires a dry run on a representative batch. A related subtlety: the random initialization now happens at first call rather than at construction, so any random numbers your program draws in between shift the generator’s state, and a fixed seed can yield different weights than the explicitly shaped version of the same model (Section 5.8 returns to seeding). NNX has no lazy mode: its linear layers require both widths and create parameters in the constructor. PyTorch’s lazy layers instead let the first real batch supply the missing width.
Until the first call, the variables do not exist. An optimizer object can be constructed without them, but reading weights, counting parameters, or building optimizer state requires the model’s variable list. Keras offers two ways to create it: a dry run on a representative batch, or net.build((None, 20)), which propagates shapes through the model without any data (None marks the batch dimension). Initialization happens at build rather than at construction, so under a fixed seed the weights you get depend on how many random numbers the program drew before the model was built (Section 5.8 returns to seeding). As the containers lesson showed, building is also the moment the model’s structure locks.
Parameters are available as soon as construction returns. Initialization randomness comes from the nnx.Rngs object passed to the constructors, so a layer’s random stream is explicit (Section 5.8 returns to seeding).
Until the first forward pass, Gluon parameters have placeholder shapes and no allocated values. A Trainer can be constructed over those parameter objects, but reading a weight or counting elements requires a dry run on a representative batch. A related subtlety: the random draw now happens at first call rather than at initialize(), so any random numbers your program draws in between shift the generator’s state, and a fixed seed can yield different weights than an explicitly shaped version of the same model (Section 5.8 returns to seeding). Other libraries retrofitted this behavior; Gluon was designed around it, which is why initialize() is a separate, explicit step rather than something the constructor does.
@d2l.add_to_class(d2l.Module)
def apply_init(self, inputs, init=None):
self.forward(*inputs)
if init is not None:
self.net.apply(init)nn.Module.apply(fn) calls fn on every module in the tree, children first. It is the standard way to push a policy across an arbitrary model, one more operation that is a tree walk, and from Chapter 6 on the idiom model.apply_init([X], init) opens most of our training scripts. A small demonstration:
Each nnx.Linear accepts a kernel_init function and invokes it in the constructor. A small demonstration:
class TinyMLP(d2l.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(),
nn.LazyLinear(10))
def init_xavier(module):
if type(module) == nn.Linear:
nn.init.xavier_uniform_(module.weight)
model = TinyMLP()
model.apply_init([X], init_xavier)
model.net[0].weight.shapetorch.Size([256, 20])net = tf.keras.Sequential([
tf.keras.layers.Dense(
256, activation='relu',
kernel_initializer=tf.keras.initializers.GlorotUniform()),
tf.keras.layers.Dense(10)])
net.build((None, 20))
net.layers[0].kernel.shapeTensorShape([20, 256])class TinyMLP(d2l.Module):
def __init__(self, rngs):
super().__init__()
self.net = nnx.Sequential(
nnx.Linear(20, 256,
kernel_init=nnx.initializers.xavier_uniform(),
rngs=rngs), nnx.relu,
nnx.Linear(256, 10, rngs=rngs))
model = TinyMLP(nnx.Rngs(0))
[(path, value.shape)
for path, value in nnx.state(model, nnx.Param).flat_state()][(('net', 'layers', 0, 'bias'), (256,)),
(('net', 'layers', 0, 'kernel'), (20, 256)),
(('net', 'layers', 2, 'bias'), (10,)),
(('net', 'layers', 2, 'kernel'), (256, 10))]net = nn.Sequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))
net.initialize(init.Xavier())
net(X)
net[0].weight.data().shape(256, 20)The dry run inside apply_init turned every lazy layer into a real one, after which init_xavier could match on nn.Linear and rewrite its weights. Which initializer to apply, and why Xavier’s variance rule (Section 4.4.2.2) is a sensible default, is the subject of Section 5.3.
build materialized every shape and drew the first kernel’s initial values from the Xavier initializer, without a batch of data in sight. Which initializer to use, and why Xavier’s variance rule (Section 4.4.2.2) is a sensible default, is the subject of Section 5.3.
The constructor created every parameter and drew the first kernel from the Xavier initializer. Which initializer to use, and why Xavier’s variance rule (Section 4.4.2.2) is a sensible default, is the subject of Section 5.3.
The dry run materialized every shape and drew the first weight’s initial values from the Xavier initializer that initialize had recorded. Which initializer to use, and why Xavier’s variance rule (Section 4.4.2.2) is a sensible default, is the subject of Section 5.3.
5.1.5 Building from a Config
So far every width in this section was a literal typed into a constructor. Real model code does not work that way. The architecture choices that vary across experiments, such as depth, width, and output size, must be logged with results and stored with checkpoints so that a saved model can be rebuilt. The topology remains in code; a small configuration object records its variable choices:
@dataclass
class MLPConfig:
d_in: int = 784
d_hidden: int = 256
num_blocks: int = 4
d_out: int = 10
def build(cfg: MLPConfig) -> nn.Module:
blocks = [ResidualBlock(cfg.d_hidden) for _ in range(cfg.num_blocks)]
return nn.Sequential(nn.Linear(cfg.d_in, cfg.d_hidden),
*blocks, nn.Linear(cfg.d_hidden, cfg.d_out))@dataclass
class MLPConfig:
d_in: int = 784
d_hidden: int = 256
num_blocks: int = 4
d_out: int = 10
def build(cfg: MLPConfig) -> tf.keras.Model:
blocks = [ResidualBlock(cfg.d_hidden) for _ in range(cfg.num_blocks)]
return tf.keras.Sequential([tf.keras.Input((cfg.d_in,)),
tf.keras.layers.Dense(cfg.d_hidden),
*blocks,
tf.keras.layers.Dense(cfg.d_out)])@dataclass
class MLPConfig:
d_in: int = 784
d_hidden: int = 256
num_blocks: int = 4
d_out: int = 10
class ResidualMLP(nnx.Module):
def __init__(self, cfg, rngs):
self.cfg = cfg
self.proj = nnx.Linear(cfg.d_in, cfg.d_hidden, rngs=rngs)
self.blocks = nnx.List([
ResidualBlock(cfg.d_hidden, rngs) for _ in range(cfg.num_blocks)])
self.out = nnx.Linear(cfg.d_hidden, cfg.d_out, rngs=rngs)
def __call__(self, X):
X = self.proj(X)
for block in self.blocks:
X = block(X)
return self.out(X)
def build(cfg: MLPConfig) -> nnx.Module:
return ResidualMLP(cfg, nnx.Rngs(0))@dataclass
class MLPConfig:
d_in: int = 784
d_hidden: int = 256
num_blocks: int = 4
d_out: int = 10
def build(cfg: MLPConfig) -> nn.Block:
net = nn.Sequential()
net.add(nn.Dense(cfg.d_hidden, in_units=cfg.d_in))
for _ in range(cfg.num_blocks):
net.add(ResidualBlock(cfg.d_hidden))
net.add(nn.Dense(cfg.d_out, in_units=cfg.d_hidden))
return netOne config produces one architecture: an input projection into the hidden width, num_blocks identical residual blocks, and an output projection. The list comprehension producing blocks is a plain Python list, which is safe here for the reason Section 5.1.2 taught: unpacking it into nn.Sequential is what registers each block. Every layer gets explicit widths, and this is where explicit shapes beat lazy ones: the config already knows every width, so explicit construction yields a fully materialized model with no dry run needed. Printing the model displays the module tree:
One config produces one architecture: an input projection into the hidden width, num_blocks identical residual blocks, and an output projection. The free-standing build(cfg) function is unrelated to the build method Keras calls on layers, but the two meet here: because the config knows the input width, we can declare it with tf.keras.Input, and Sequential then builds the whole model at construction time, no dry run needed. This is where explicit shapes beat inferred ones. net.summary() displays the module tree, every shape and parameter count already known:
One config produces one architecture. The input width is explicit because NNX creates parameters in the constructor. nnx.List registers the repeated residual blocks, while the loop in __call__ controls their execution. Selecting parameter state displays the resulting object graph:
One config produces one architecture: an input projection into the hidden width, num_blocks identical residual blocks, and an output projection. This is where explicit shapes beat deferred ones: the config already knows every width, so the two projections take in_units directly and skip deferred initialization. (The Dense layers inside each ResidualBlock were written without in_units, so the model as a whole still wants one dry run before its parameters can be counted.) Printing the model displays the module tree:
net = build(MLPConfig())
netSequential(
(0): Linear(in_features=784, out_features=256, bias=True)
(1): ResidualBlock(
(body): Sequential(
(0): Linear(in_features=256, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=256, bias=True)
)
)
(2): ResidualBlock(
(body): Sequential(
(0): Linear(in_features=256, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=256, bias=True)
)
)
(3): ResidualBlock(
(body): Sequential(
(0): Linear(in_features=256, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=256, bias=True)
)
)
(4): ResidualBlock(
(body): Sequential(
(0): Linear(in_features=256, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=256, bias=True)
)
)
(5): Linear(in_features=256, out_features=10, bias=True)
)net = build(MLPConfig())
net.summary()Model: "sequential_9"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type) ┃ Output Shape ┃ Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ dense_25 (Dense) │ (None, 256) │ 200,960 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ residual_block_1 │ (None, 256) │ 131,584 │
│ (ResidualBlock) │ │ │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ residual_block_2 │ (None, 256) │ 131,584 │
│ (ResidualBlock) │ │ │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ residual_block_3 │ (None, 256) │ 131,584 │
│ (ResidualBlock) │ │ │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ residual_block_4 │ (None, 256) │ 131,584 │
│ (ResidualBlock) │ │ │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_26 (Dense) │ (None, 10) │ 2,570 │
└─────────────────────────────────┴────────────────────────┴───────────────┘ Total params: 729,866 (2.78 MB) Trainable params: 729,866 (2.78 MB) Non-trainable params: 0 (0.00 B)net = build(MLPConfig())
[(path, value.shape)
for path, value in nnx.state(net, nnx.Param).flat_state()][(('blocks', 0, 'body', 'layers', 0, 'bias'), (256,)),
(('blocks', 0, 'body', 'layers', 0, 'kernel'), (256, 256)),
(('blocks', 0, 'body', 'layers', 2, 'bias'), (256,)),
(('blocks', 0, 'body', 'layers', 2, 'kernel'), (256, 256)),
(('blocks', 1, 'body', 'layers', 0, 'bias'), (256,)),
(('blocks', 1, 'body', 'layers', 0, 'kernel'), (256, 256)),
(('blocks', 1, 'body', 'layers', 2, 'bias'), (256,)),
(('blocks', 1, 'body', 'layers', 2, 'kernel'), (256, 256)),
(('blocks', 2, 'body', 'layers', 0, 'bias'), (256,)),
(('blocks', 2, 'body', 'layers', 0, 'kernel'), (256, 256)),
(('blocks', 2, 'body', 'layers', 2, 'bias'), (256,)),
(('blocks', 2, 'body', 'layers', 2, 'kernel'), (256, 256)),
(('blocks', 3, 'body', 'layers', 0, 'bias'), (256,)),
(('blocks', 3, 'body', 'layers', 0, 'kernel'), (256, 256)),
(('blocks', 3, 'body', 'layers', 2, 'bias'), (256,)),
(('blocks', 3, 'body', 'layers', 2, 'kernel'), (256, 256)),
(('out', 'bias'), (10,)),
(('out', 'kernel'), (256, 10)),
(('proj', 'bias'), (256,)),
(('proj', 'kernel'), (784, 256))]net = build(MLPConfig())
net.initialize()
netSequential(
(0): Dense(784 -> 256, linear)
(1): ResidualBlock(
(body): Sequential(
(0): Dense(-1 -> 256, Activation(relu))
(1): Dense(-1 -> 256, linear)
)
)
(2): ResidualBlock(
(body): Sequential(
(0): Dense(-1 -> 256, Activation(relu))
(1): Dense(-1 -> 256, linear)
)
)
(3): ResidualBlock(
(body): Sequential(
(0): Dense(-1 -> 256, Activation(relu))
(1): Dense(-1 -> 256, linear)
)
)
(4): ResidualBlock(
(body): Sequential(
(0): Dense(-1 -> 256, Activation(relu))
(1): Dense(-1 -> 256, linear)
)
)
(5): Dense(256 -> 10, linear)
)net(torch.rand(2, 784)).shapetorch.Size([2, 10])net(tf.random.uniform((2, 784))).shapeTensorShape([2, 10])net(jax.random.uniform(d2l.get_key(), (2, 784))).shape(2, 10)net(np.random.uniform(size=(2, 784))).shape(2, 10)The variable architecture choices are now data. Rescaling the model is a change to two fields, not an edit to model code:
for cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):
n = sum(p.numel() for p in build(cfg).parameters())
print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '
f'{n:,} parameters')d_hidden=256, num_blocks=4: 729,866 parameters
d_hidden=512, num_blocks=8: 4,609,546 parametersfor cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):
print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '
f'{build(cfg).count_params():,} parameters')d_hidden=256, num_blocks=4: 729,866 parametersd_hidden=512, num_blocks=8: 4,609,546 parametersfor cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):
net = build(cfg)
n = sum(x.size for _, x in nnx.state(net, nnx.Param).flat_state())
print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '
f'{n:,} parameters')d_hidden=256, num_blocks=4: 729,866 parametersd_hidden=512, num_blocks=8: 4,609,546 parametersfor cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):
net = build(cfg)
net.initialize()
net(np.zeros((2, cfg.d_in))) # Dry run: materialize the deferred shapes
n = sum(p.data().size for p in net.collect_params().values())
print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '
f'{n:,} parameters')d_hidden=256, num_blocks=4: 729,866 parameters
d_hidden=512, num_blocks=8: 4,609,546 parametersBecause build is deterministic in cfg, the config is all you need to reconstruct the module tree later; Section 5.6 saves it alongside the weights so that loading a checkpoint starts by rebuilding the exact same model. A config of widths and depths feeding a loop that stacks identical residual blocks is a common construction pattern in the Transformer implementations later in the book.
Because build is deterministic in cfg, the config is all you need to reconstruct the module tree later; Section 5.6 saves it alongside the weights so that loading a checkpoint starts by rebuilding the exact same model. Keras bakes the same idea into every layer: get_config() returns the constructor arguments needed to re-create the object, and that is what Keras model serialization records. A config of widths and depths feeding a loop that stacks identical residual blocks is a common construction pattern in the Transformer implementations later in the book.
Because build is deterministic in cfg, the config is all we need to reconstruct the object graph later; Section 5.6 saves it alongside the state. Width and depth fields feeding a loop that stacks identical residual blocks form a common construction pattern in the Transformer models later in the book.
Because build is deterministic in cfg, the config is all you need to reconstruct the module tree later; Section 5.6 saves it alongside the weights so that loading a checkpoint starts by rebuilding the exact same model. A config of widths and depths feeding a loop that stacks identical residual blocks is a common construction pattern in the Transformer implementations later in the book.
5.1.6 Summary
A module owns parameters, child modules, and a forward method. Layers, blocks, and whole models are the same kind of object, so the usual hierarchy is tree-shaped, while shared children introduce aliases. Parameter collection, device movement, and serialization traverse that object graph. Children are discovered through registration: attribute assignment and the containers nn.Sequential, nn.ModuleList, and nn.ModuleDict register children, while a plain Python list hides them, yielding a model that runs but trains nothing. forward is ordinary Python; a residual connection is one line in it. Lazy layers declare output widths and infer input widths on the first forward pass, so initialization and inspection follow a dry run (apply_init). Configs turn architecture into data: a dataclass of widths and depths plus a build function that stacks blocks.
A module owns variables, child layers, and a call method. Layers, blocks, and whole models are the same kind of object, so the usual hierarchy is tree-shaped, while shared children introduce aliases. Variable collection and serialization traverse that object graph. Children are discovered through attribute assignment: Keras scans every assigned attribute, lists and dictionaries included, so even a plain Python list of layers is tracked. Variables are created by build, not by the constructor: every layer infers its input width when the first batch arrives (or when build(input_shape) is called), after which the model’s structure is locked. call is ordinary Python; a residual connection is one line in it. Configs turn architecture into data: a dataclass of widths and depths plus a build(cfg) function that stacks blocks, with get_config() as the Keras-native form of the same idea.
A module owns parameters, child modules, and a __call__ method. Layers, blocks, and whole models are the same kind of object, and state selection, optimizer updates, and serialization traverse the NNX object graph. nnx.List and nnx.Dict register collections of children; a plain container holding graph nodes is rejected. __call__ is ordinary Python, so a residual connection is one line. NNX linear layers take both widths and create their parameters in the constructor. A dataclass config plus a deterministic build function records those widths and repeated-block counts.
A module owns parameters, child blocks, and a forward method. Layers, blocks, and whole models are the same kind of object, so the usual hierarchy is tree-shaped, while shared children introduce aliases. Parameter collection, initialization, device movement, and serialization traverse that object graph. Children are discovered through registration: attribute assignment, nn.Sequential’s add, and register_child register children, while a plain Python list hides them, though collect_params() warns about the hidden blocks by name. forward is ordinary Python; a residual connection is one line in it. Deferred initialization is Gluon’s native mode: every layer declares its output width, initialize() records the initializer, and the first forward pass infers input widths and allocates the parameters, so inspection and training follow a dry run. Configs turn architecture into data: a dataclass of widths and depths plus a build function that stacks blocks.
5.1.7 Exercises
- Take
PlainListMLPand catalog everything that breaks besides the empty parameter list. Checknet.state_dict(), the effect ofnet.to(torch.float64)on the hidden layers’ dtypes, and whethernet.eval()reaches the children. Explain how each failure follows from the same missing registration. - Implement a
ParallelBlockthat takes two child modulesnet1andnet2, runs both on the same input, and concatenates their outputs along the last dimension. What must be true of the two children’s outputs for the concatenation to be valid? - Extend
MLPConfigwith an activation switch (for example,act: str = 'relu') and makebuildhonor it. Which decisions belong in a config and which belong in code? Where would you put a choice betweenResidualBlockand a plain feed-forward block? ResidualBlockrequires its input and output widths to agree. Suppose you want a block whose output is wider than its input. Give two standard fixes and the cost of each. (Section 7.4 uses one of them in ResNet.)
- Keras’s tracking has one blind spot left: create a
Denselayer insidecallrather than in the constructor. The model runs without complaint; checklen(net.trainable_variables)after calling it, explain what an optimizer would train, and explain what happens to the layer’s weights between two calls. - Implement a
ParallelBlockthat takes two child modulesnet1andnet2, runs both on the same input, and concatenates their outputs along the last dimension. What must be true of the two children’s outputs for the concatenation to be valid? - Extend
MLPConfigwith an activation switch (for example,act: str = 'relu') and makebuildhonor it. Which decisions belong in a config and which belong in code? Where would you put a choice betweenResidualBlockand a plain feed-forward block? ResidualBlockrequires its input and output widths to agree. Suppose you want a block whose output is wider than its input. Give two standard fixes and the cost of each. (Section 7.4 uses one of them in ResNet.)
- Extend
ListMLPwith annnx.Dictof submodules and verify that every parameter appears innnx.state(net, nnx.Param). Then replace it with a plain dictionary holding those modules and describe how NNX reports the invalid graph container. - Implement a
ParallelBlockthat takes two child modulesnet1andnet2as fields, runs both on the same input, and concatenates their outputs along the last dimension. What must be true of the two children’s outputs for the concatenation to be valid? - Extend
MLPConfigwith an activation switch (for example,act: str = 'relu') and makeResidualMLPhonor it. Which decisions belong in the config and which belong in code? Where would you put a choice betweenResidualBlockand a plain feed-forward block? ResidualBlockrequires its input and output widths to agree. Suppose you want a block whose output is wider than its input. Give two standard fixes and the cost of each. (Section 7.4 uses one of them in ResNet.)
- Take
PlainListMLPand catalog everything that breaks besides the empty parameter dictionary. Checknet.save_parameters(...), whatnet.initialize()reaches, and what happens on the first forward pass if you remove the by-hand initialization from the constructor. Explain how each failure follows from the same missing registration, and note where along the way thecollect_params()warning fires. - Implement a
ParallelBlockthat takes two child blocksnet1andnet2, runs both on the same input, and concatenates their outputs along the last dimension. What must be true of the two children’s outputs for the concatenation to be valid? - Extend
MLPConfigwith an activation switch (for example,act: str = 'relu') and makebuildhonor it. Which decisions belong in a config and which belong in code? Where would you put a choice betweenResidualBlockand a plain feed-forward block? ResidualBlockrequires its input and output widths to agree. Suppose you want a block whose output is wider than its input. Give two standard fixes and the cost of each. (Section 7.4 uses one of them in ResNet.)