Batching And Parameter Sharing🔗

Consider an ordinary one-hidden-layer network

f_\theta(x)=W_2\,\rho(W_1x+b_1)+b_2.

For a batch of B vectors the shapes are

X\in\mathbb{R}^{B\times d_{\mathrm{in}}},\qquad W_1\in\mathbb{R}^{d_h\times d_{\mathrm{in}}},\qquad W_2\in\mathbb{R}^{d_{\mathrm{out}}\times d_h}.

In TorchLean the per-sample model and the batch axis are written separately. The per-sample model maps [3] to [2], and nn.mapLeading applies it at every index of a new leading axis, sharing one parameter set.

The distinction becomes useful when changing batch size. Eight examples require eight sets of activations, but there is still only one classifier to update. In this example the first affine map contributes fifteen weights and five biases; the second contributes ten weights and two biases. Those thirty-two scalars are shared by every row. The printed shapes let us account for them without inspecting an optimizer or relying on a naming convention.

[[5, 3], [5], [2, 5], [2]]

The parameter list is the per-sample list: adding a batch axis added no parameters. Running the batched model on eight copies of the same row gives eight copies of the same output row, which is what parameter sharing means operationally.

Evaluation mode also makes this a controlled comparison: the model contains affine maps and a ReLU, with no operation that depends on other batch members. Equal rows therefore take equal paths through the same parameters. The values themselves are determined by the seeded initial model; they are predictions before training, not class probabilities or evidence that a task has been learned.

[[0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478],
 [0.428865, 0.706478]]

PyTorch writes the same model without mentioning the batch axis at all, because nn.Linear maps over any number of leading dimensions:

# Inspect the shared weights, then compare all output rows
# with the first.
mlp = nn.Sequential(nn.Linear(3, 5), nn.ReLU(), nn.Linear(5, 2))
print([tuple(p.shape) for p in mlp.parameters()])
out = mlp(torch.ones(8, 3))
print(tuple(out.shape), bool(torch.all(out == out[0])))

[(5, 3), (5,), (2, 5), (2,)]
(8, 2) True

Both versions share parameters across the batch. Because [8, 3] appears in the TorchLean model's type, a theorem or an exported graph about this model states which batch shape it is about, and an operation that is not batch-invariant cannot hide behind an implicit leading axis; TorchLean and PyTorch shows a real batched-versus-sliced discrepancy of that kind. The cost is that the batch axis has to be named, and nn.mapLeading is where you name it.

The batch model fixes how these affine maps and activations compose. Adding a residual branch introduces another requirement: the two paths must agree on the shape at their join.

Residual Networks🔗

A residual block computes

y=F_\theta(x)+S_\theta(x),

where S is either the identity or a projection (He et al., 2016). The addition is defined only when the two branches produce the same shape, so nn.addBranches requires exactly that, and the parameter list of the block is the concatenation of the two branch lists:

[[8, 32], [8], [8, 32], [8]]

Give the shortcut the wrong width and there is no model to run:

Application type mismatch: The argument
  shortcut
has type
  nn.Sequential (Shape.appendDim [] 32) (Shape.appendDim [] 6)
but is expected to have type
  nn.Sequential [32] [8]
in the application
  nn.addBranches wide shortcut

With the identity as the shortcut and a single ReLU as the branch, the block has no parameters. It computes x+\rho(x), so a negative coordinate passes through unchanged and a positive one is doubled:

x + relu x = [-1.000000, 4.000000]

# Use the same two coordinates to expose the sign-dependent
# residual behavior.
x = torch.tensor([-1.0, 2.0])
print((x + torch.relu(x)).tolist())

[-1.0, 4.0]

The two residual examples separate two choices that are easy to conflate. mmBranches learns a new representation on both paths because its input and output widths differ. mmRes keeps the input representation and adds a correction of the same width. The latter makes the identity path visible in the negative coordinate: even when ReLU contributes zero, the original value survives. For a learned residual branch the same addition permits a small correction without requiring that branch to reproduce the whole input.

-- Spatial rank controls the length of the grid and
-- kernel-radius vectors.
structure ResNet.Config (d : Nat) where
  inputChannels  : Nat
  spatial        : Tensor Nat [d]
  hiddenChannels : Nat
  kernelRadius   : Tensor Nat [d] := Tensor.ones [d]
  classCount     : Nat

([8, 3, 32, 32], [8, 16, 32, 32], [8, 10])

@nn.Convolution.Geometry.outputSpatial_samePadding : ∀ {d : ℕ} (input radius : Tensor ℕ [d]),
  (nn.Convolution.Geometry.samePadding radius).outputSpatial input = input

# Pool spatial positions only; the channel dimension becomes
# the classifier input.
conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
images = torch.zeros(8, 3, 32, 32)
print(tuple(conv(images).shape))
pooled = F.adaptive_avg_pool2d(conv(images), 1).flatten(1)
print(tuple(pooled.shape), tuple(nn.Linear(16, 10)(pooled).shape))

(8, 16, 32, 32)
(8, 16) (8, 10)

Vision Transformers🔗

A vision Transformer first turns a spatial field into a token sequence (Dosovitskiy et al., 2021). For an input with spatial extent n_1\times\cdots\times n_d, patch kernel k, stride s, and padding p, each output extent is the usual convolution expression

n'_i =\left\lfloor\frac{n_i+2p_i-k_i}{s_i}\right\rfloor+1.

If the patch convolution emits D channels then the patch grid becomes

B\times D\times n'_1\times\cdots\times n'_d \;\longrightarrow\; B\times N\times D,\qquad N=\prod_i n'_i.

NN.API.Models.Vit defines that conversion as patchesToTokens: the implementation reshapes the patch grid and moves the channel axis to the end, so the following Transformer block receives the conventional batch × sequence × embedding layout. Every stage is a derived field of the configuration, so we can print the whole pipeline. Take 32\times32 inputs with 4\times4 non-overlapping patches, sixty-four embedding channels, four heads of width sixteen, two blocks, and a learned class slot:

([8, 8], 64, 65)

A 32\times32 image with stride-four patches has an 8\times8 patch grid, hence sixty-four tokens, and the class slot makes the encoded sequence length sixty-five. The three intermediate tensor shapes follow, for a batch of eight:

([8, 64, 8, 8], [8, 64, 64], [8, 65, 64])

([8, 3, 32, 32], [8, 10])

In [8, 64, 64], the two axes of length sixty-four mean different things: one counts patches and one counts features within a patch token. Their equal sizes happen to hide a possible transpose mistake. The preceding [8, 64, 8, 8] and following [8, 65, 64] disambiguate them: adding the class slot changes the sequence length, while the embedding width stays fixed. The endpoint [8, 10] then says that the classifier produces one score vector per image, rather than one per patch.

The corresponding PyTorch operations expose the reshape and axis permutation:

# Flatten only the patch grid, then move channels behind the
# sequence axis.
patch = nn.Conv2d(3, 64, kernel_size=4, stride=4)
grid = patch(torch.zeros(8, 3, 32, 32))
print(tuple(grid.shape))
tokens = grid.flatten(2).transpose(1, 2)
print(tuple(tokens.shape))
cls = torch.zeros(8, 1, 64)
print(tuple(torch.cat([cls, tokens], dim=1).shape))

(8, 64, 8, 8)
(8, 64, 64)
(8, 65, 64)

The intermediate shapes agree at every stage. flatten(2).transpose(1, 2) computes the token tensor at runtime, while mmVit.encoder.patchTokenShape [8] computes its shape without running the model. Computing that shape does not run validation; the model constructor separately calls ViT.EncoderConfig.validate to reject an unusable patch grid.

([6, 6], 36, 37)

# A stride-five convolution emits only complete patches from
# the 32-pixel grid.
patch = nn.Conv2d(3, 64, kernel_size=5, stride=5)
print(tuple(patch(torch.zeros(8, 3, 32, 32)).shape))

(8, 64, 6, 6)

Recurrence, Attention, And Causality🔗

Sequence models add state or a causal dependency. Write h_{t-1} for the state before input x_t and h_t for the updated state. A vanilla recurrent network then computes

h_t=\phi(W_xx_t+W_hh_{t-1}+b),\qquad y_t=W_yh_t+b_y.

An LSTM replaces the single update by input, forget, output, and candidate gates:

\begin{aligned} i_t&=\sigma(W_ix_t+U_ih_{t-1}+b_i),\\ f_t&=\sigma(W_fx_t+U_fh_{t-1}+b_f),\\ o_t&=\sigma(W_ox_t+U_oh_{t-1}+b_o),\\ \tilde c_t&=\tanh(W_cx_t+U_ch_{t-1}+b_c),\\ c_t&=f_t\odot c_{t-1}+i_t\odot\tilde c_t,\\ h_t&=o_t\odot\tanh(c_t). \end{aligned}

In the cell update, f_t scales the old cell state and i_t scales the candidate \tilde c_t. The output gate o_t controls how much of the transformed cell state appears in h_t. These elementwise products require each gate to have the same width as the state it acts on.

The hidden and cell states are explicit tensors whose dimensions have to stay stable across the unrolled sequence. The runnable rnn, lstm, and lstm_regression commands use that typed state path. There is currently no GRU training subcommand: the LiRPA verifier has a GRU-gate certificate fixture, but that is a different artifact and should not be presented as a trainable GRU model.

An explicit state gives a recurrent model a streaming interface: a caller can supply the state left by a previous chunk and continue the computation. Restarting from zero instead describes a new sequence. For an LSTM there are two pieces to retain, the hidden state and the cell state; keeping only the hidden output loses information needed by the next update. Their matching widths make the gate products legal, but correct streaming also requires carrying the actual values.

[[true, false, false, false], [true, true, false, false], [true, true, true, false], [true, true, true, true]]

[[1.000000, 2.000000], [2.339523, 3.339523], [3.510470, 4.510470]]

[[3.000000, 4.000000], [3.406673, 4.406673], [3.510470, 4.510470]]

# Use binary64 to compare the same causal and unrestricted
# weighted sums.
q = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=torch.float64)
v = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float64)
print(F.scaled_dot_product_attention(q, q, v, is_causal=True).tolist())
print(F.scaled_dot_product_attention(q, q, v).tolist())

[[1.0, 2.0], [2.3395230986533138, 3.3395230986533138],
 [3.5104695304536615, 4.510469530453662]]
[[3.0, 4.0], [3.4066725560787154, 4.406672556078716],
 [3.5104695304536615, 4.510469530453662]]

[[1.000000, 2.000000], [2.339523, 3.339523], [3.510470, 4.510470], [5.837654, 6.837654]]

[[3.000000, 4.000000], [3.406673, 4.406673], [3.510470, 4.510470]]

[[5.000000, 6.000000], [4.000000, 5.000000], [3.510470, 4.510470]]

# Compare unrestricted permutation symmetry with the
# position-dependent causal mask.
q = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=torch.float64)
k = torch.tensor([[1.0, 1.0], [0.0, 1.0], [1.0, 0.0]], dtype=torch.float64)
v = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float64)
vp = torch.tensor([[5.0, 6.0], [3.0, 4.0], [1.0, 2.0]], dtype=torch.float64)
print(F.scaled_dot_product_attention(q, q, v).tolist())
print(F.scaled_dot_product_attention(q, k, vp).tolist())
print(F.scaled_dot_product_attention(q, k, vp, is_causal=True).tolist())

[[3.0, 4.0], [3.4066725560787154, 4.406672556078716],
 [3.5104695304536615, 4.510469530453662]]
[[3.0, 4.0], [3.4066725560787154, 4.406672556078715],
 [3.510469530453662, 4.510469530453662]]
[[5.0, 6.0], [4.0, 5.0],
 [3.510469530453662, 4.510469530453662]]

0.000000

@hardMaskedSoftmaxSpec_causal_future_zero : ∀ {n : ℕ} (scores : Tensor ℝ [n, n]) (i j : Fin n),
  ↑i < ↑j → Spec.get2 (Spec.hardMaskedSoftmaxSpec scores (Spec.causalMask n)) i j = 0

token embedding
  -> learned positional embedding
  -> masked Transformer blocks
  -> layer normalization
  -> vocabulary projection

State-Space Models🔗

Mamba-style models return to explicit state but allow input-dependent scan parameters (Gu and Dao, 2024). Using the same convention of reading the updated state, a simplified state-space recurrence is

h_t=A_t h_{t-1}+B_t x_t,\qquad y_t=C_t h_t+D_t x_t.

NN.API.Models.Mamba holds the trainable side of this:

• Mamba.Config, which records the vocabulary size and the model width, and whose validate rejects a zero-length sequence before any tensor is allocated;

• Mamba.languageModel, a trainable selective state-space core followed by a linear projection back to vocabulary logits at every time step. The core is built from autograd-covered operations, so --device cuda trains the same parameters on the CUDA backend.

Validation is an explicit function on the configuration and sequence length. Calling it before model construction rejects invalid extents before allocating model tensors:

Except.error "Mamba: sequence length must be positive"

Except.ok ()

([8, 256], [8, 256])

[0, 256] is a valid tensor shape, but this model requires a positive sequence length. validate checks that additional requirement and returns an Except.error before a scan runs. The successful call with length eight separates this configuration check from the input and output shape calculation printed below it.

Input and output shapes coincide here because the model predicts a distribution over the same vocabulary at every position, which is what makes next-token training a matter of shifting the target by one rather than of reshaping anything.

Here the input rows are one-hot vocabulary features and the output rows are logits. Equal shapes do not mean equal representations: the input selects a token category, while a logit row scores all possible next categories. A loss or sampling rule interprets those scores. The hidden modelWidth can therefore be thirty-two even though the public input and output width is 256; the vocabulary projection restores the external width after the recurrent computation.

The trainable core computes \Delta, B_t and C_t from the current token features. The specification describes the corresponding input-dependent recurrence as SelectiveMambaBlockSpec in NN.Spec.Models.Mamba, with runArray and runArrayWithHistory runners. The API constructor supplies trainable parameters; the spec runners expose the state and history needed to state properties of the recurrence.

The causality property proved for attention above has a recurrent counterpart, and it is stated for the array runners in MambaCausality.lean: for any initial state h_0 and any sequences xs and ys,

\operatorname{take}\,|xs|\;\bigl(\operatorname{run}(h_0,\;xs\mathbin{+\!+}ys)\bigr) =\operatorname{run}(h_0,\;xs).

Appending future tokens cannot change outputs already emitted for a prefix. The file proves this for the diagonal S4 runner, the compact Mamba block, and the full selective block, including the variant that carries a causal-convolution history, all on top of the scan algebra in Scan.lean. Stating it at the runner level is deliberate: a runtime may use a chunked or parallel selective scan, but it has to refine these runners.

The trainable constructor currently uses a selective Mamba-1 block with explicit convolution history and input-dependent state parameters. Its time-step projection is stored as a dense matrix. Loading factors from a low-rank checkpoint would require multiplying those factors to match this forward parameterization; training the resulting dense matrix would then expose a different set of independent parameters. The compact block below is chosen to make the recurrence arithmetic visible, and should be read as that smaller spec example.

#[[1.500000], [-0.250000], [1.875000]]

#[[1.500000], [-0.250000], [1.875000], [2.687500], [-1.531250]]

true

Neural Operators🔗

An FNO learns a map between functions sampled on a grid rather than a map between fixed feature vectors (Li et al., 2021). One spectral block has the schematic form

v_{\ell+1}(x) =\sigma\!\left(W_\ell v_\ell(x) +\mathcal F^{-1}\!\left(R_\ell\cdot\mathcal F(v_\ell)\right)(x)\right).

The two terms combine local and grid-wide information. W_\ell mixes channels at each grid point. The spectral term transforms the whole field, applies learned weights R_\ell to retained Fourier coefficients, and transforms back. The activation \sigma acts on the sum. The retained mode set therefore determines which frequencies the spectral branch can change.

The configuration in NN.API.Models.FNO is again parameterized by spatial rank, and the field-to-field boundary is visible in its shapes: a grid of sixty-four points with a mode-band width of twelve, a latent width of thirty-two, and four blocks maps a sampled field to a sampled field of the same extent.

([8, 64], [8, 64])

On each full-DFT axis, modes is the width of each of two index bands. A width of four on a sixteen-point axis therefore retains eight indices. The second example uses an eight-point axis whose bands overlap:

([0, 1, 2, 3, 12, 13, 14, 15], [0, 1, 2, 3, 4, 5, 6, 7])

For an axis of length 16, indices 0,\ldots,3 represent the zero and low positive frequencies; indices 12,\ldots,15 represent frequencies -4,\ldots,-1. These are two retained bands, not exact positive/negative pairs: the mask includes -4 but excludes +4. When the bands overlap, the whole axis is retained without double counting. The PyTorch transcript uses the same FFT index convention:

# Translate storage indices into signed frequencies before
# selecting the two bands.
freqs = (torch.fft.fftfreq(16) * 16).to(torch.int64)
print(freqs.tolist())
print([i for i in range(16) if abs(int(freqs[i])) < 4 or int(freqs[i]) == -4])

[0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1]
[0, 1, 2, 3, 12, 13, 14, 15]

The frequency lists give a more precise description than “keep four modes.” The first list maps every storage location to a signed frequency; the second selects the locations the layer actually uses. Comparing them reveals the asymmetric endpoint at negative four. It also explains why asking for six indices at each end of an eight-point axis keeps everything: the two sets cover the axis. This is set membership, so their overlap does not multiply a coefficient twice.

The public constructor transforms each spatial axis separately. With spectralPath := .automatic, eager CUDA execution uses cuFFT and other interpreters use dense per-axis operations. Setting spectralPath := .denseReference selects full-grid DFT matrices without changing the weights or checkpoint layout. Both paths compute the same full-spectrum model, up to floating-point differences in their transform algorithms.

The CUDA Burgers command deliberately selects a separate one-dimensional real-FFT parameterization backed by cuFFT. It has the same typed field-to-field boundary, but its one-sided weights are not interchangeable with the full-spectrum model's weights. The spectral-block chapter in Scientific Forward Models says which parts are shared.

The field shape [8,64] records eight independent sampled fields, each with sixty-four spatial values. It does not describe the hidden channel width, the retained Fourier bands, or how samples are spaced physically. Those choices live in the model and data configuration. Two operators can therefore share this boundary while learning different maps. When comparing their predictions, match the grid and coordinate convention as well as the tensor extents; a shifted field can have the right number of entries while representing a different spatial function.

References🔗

The architectures are He et al. (2016), Vaswani et al. (2017), Dosovitskiy et al. (2021), Gu and Dao (2024), and Li et al. (2021). The framework whose shapes and numbers we compare against throughout is Paszke et al. (2019).

• PyTorch pages for the operations used in the transcripts: nn.Conv2d, torch.fft, and the entry for scaled_dot_product_attention on the torch.nn.functional page.

• TorchLean sources shown above: Blocks, Attention, CausalMask, and Fno.