Chains compose: reshape, project, slice, reduce, return¶
Single ops are table stakes — the point is a whole block staying checked end to end. Ten tracked steps, one signature, no intermediate annotations; every middle shape is inferred and re-checked:
import numpy as np
from numty import F64, batch, seq
def autoencode(
patches: F64[batch, seq, 16],
w_enc: F64[16, 8],
w_dec: F64[8, 16],
) -> F64[batch, seq, 16]:
h = np.matmul(patches, w_enc) # (B, S, 8)
h = np.maximum(h, 0.0) # relu
z = h.reshape(h.shape[0], h.shape[1], 2, 4) # (B, S, 2, 4)
zt = z.transpose(0, 1, 3, 2) # (B, S, 4, 2)
f = zt.reshape(zt.shape[0], zt.shape[1], 8) # (B, S, 8)
y = np.matmul(f, w_dec) # (B, S, 16)
first = y[:, 0, :] # (B, 16)
mu = y.mean(axis=1) # (B, 16)
gate = first * mu # (B, 16)
out = gate[:, None, :] + patches # (B, S, 16)
return out
Two things the chain exercises deliberately: reshapes spell their
kept dims through .shape (bare TypeVars are not runtime values —
h.reshape(batch, seq, 2, 4) would crash at import, and the checker
rejects it too), and rebinding a new shape to an old name needs
allow_redefinition = True (recommended) — under default mypy
it is an [assignment] error, not a silent widening. This exact
function is tests/fixtures/chained_block.py: checked clean, runs to
(5, 7, 16).