Modules and parameters

Build reusable layers and residual blocks, inspect registered variables, and understand shared parameter ownership.

These outputs were saved from local execution. The documentation renderer does not execute notebook cells.

Show code cell source

Hide code cell source

from pathlib import Path
import sys
ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import jax
import jax.numpy as jnp
import numpy as np
import jaxsaw as js
from jaxsaw import nn, optim, data, checkpoint, inspection
from jaxsaw.nn import functional as F
print(f"JAXsaw {js.__version__} · JAX {jax.__version__}")
print(jax.devices())

JAXsaw 0.1.0 · JAX 0.5.3
[CudaDevice(id=0)]

Define a residual module

A module owns computation and state. The residual addition is ordinary array computation; it does not require a special model class.

class ResidualBlock(nn.Module):
    def __init__(self, width, *, rngs):
        super().__init__()
        self.hidden = nn.Linear(width, width, rngs=rngs)
        self.output = nn.Linear(width, width, rngs=rngs)

    def forward(self, x):
        return x + self.output(F.relu(self.hidden(x)))

model = nn.Sequential(nn.Linear(3, 8, rngs=js.Rngs(params=1)),
                      ResidualBlock(8, rngs=js.Rngs(params=2)),
                      nn.Linear(8, 2, rngs=js.Rngs(params=3)))
x = jnp.ones((4, 3))
print('Output shape:', model(x).shape)
print(inspection.summary(model))

Output shape: (4, 2)
Module | Type | Own parameters
items.0 | Linear | 32
items.1.hidden | Linear | 72
items.1.output | Linear | 72
items.1 | ResidualBlock | 0
items.2 | Linear | 18
<root> | Sequential | 0
Total unique parameters: 194

Inspect parameter paths

Parameters are live variables. Use .value for their arrays. Stable paths allow inspection and strict state loading.

for name, parameter in model.named_parameters():
    print(f'{name:35s} {str(parameter.value.shape):12s} {parameter.value.dtype}')
snapshot = model.state_dict()
print('Snapshot arrays:', len(snapshot))

items.0.bias                        (8,)         float32
items.0.weight                      (8, 3)       float32
items.1.hidden.bias                 (8,)         float32
items.1.hidden.weight               (8, 8)       float32
items.1.output.bias                 (8,)         float32
items.1.output.weight               (8, 8)       float32
items.2.bias                        (2,)         float32
items.2.weight                      (2, 8)       float32
Snapshot arrays: 8

Share weights deliberately

Reusing the same module ties weights. The optimizer sees each parameter once, and gradients from every invocation accumulate into that shared variable.

shared = nn.Linear(8, 8, rngs=js.Rngs(params=4))
tied = nn.Sequential(shared, nn.Tanh(), shared)
print('Unique parameter arrays:', len(list(tied.parameters())))
print(inspection.describe(tied)['root']['fields']['items'][2])
loss, gradients = js.value_and_grad(lambda m: jnp.square(m(jnp.ones((2,8)))).mean())(tied)
optimizer = optim.SGD(tied.parameters(), lr=.01)
optimizer.step(gradients)
print('Shared-layer objective:', float(loss))

Unique parameter arrays: 2
{'alias': 'model.items.0'}
Shared-layer objective: 0.023256195709109306

Return structured predictions

Multiple outputs can be a dictionary or tuple. Keep the output structure fixed inside a compiled signature. The objective decides which outputs contribute to loss.

class EncoderWithHead(nn.Module):
    def __init__(self, *, rngs):
        super().__init__()
        self.encoder = nn.Linear(3, 4, rngs=rngs)
        self.head = nn.Linear(4, 2, rngs=rngs)
    def forward(self, x):
        latent = F.relu(self.encoder(x))
        return {'latent': latent, 'logits': self.head(latent)}
output = EncoderWithHead(rngs=js.Rngs(params=5))(x)
print({name: array.shape for name,array in output.items()})

{'latent': (4, 4), 'logits': (4, 2)}