Combinators and models

StableTested and in active use. Breaking changes will be called out.

Combinators are modules that take modules. Because composition happens in the type system, a network is built by nesting types rather than by appending to a list at runtime.

from noeira.nn.combinators.sequential import Sequential
from noeira.nn.combinators.residual import Residual
from noeira.nn.combinators.repeat import Repeat
from noeira.nn.combinators.parallel import Parallel

Papers

• Deep Residual Learning for Image RecognitionHe, Zhang, Ren, Sun · 2015 · CVPR 2016 · arXiv:1512.03385ResNet.
• Attention Is All You NeedVaswani et al. · 2017 · NeurIPS 2017 · arXiv:1706.03762The transformer.
• An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleDosovitskiy et al. · 2020 · ICLR 2021 · arXiv:2010.11929ViT.

The core four

Sequential[A, B, C, …] — chain modules. Output dimensions must match the next input, checked at compile time.

comptime Net = Sequential[
    LinearReLU[784, 256],
    LinearReLU[256, 128],
    Linear[128, 10],
]

Residual[M] — add the input to the output. Requires IN_DIM == OUT_DIM.

comptime Block = Residual[Sequential[Linear[D, D], ReLU[D], Linear[D, D]]]

Repeat[M, N] — N independent copies in sequence. Each gets its own parameters and its own random stream (a Dropout in the block draws a different mask in each copy); the count is compile-time, so there is no runtime loop.

comptime Trunk = Repeat[6, Block]      # six residual blocks

Parallel[A, B, …] — run several modules on the same input and combine. This is where ARITY earns its place: the combinator handles children with different input counts uniformly.

The rest

Combinator	Purpose
BranchConcat	branch, then concatenate outputs
SkipConcat	concatenate a skip connection — U-net style
ProjectedResidual	residual with a projection when dimensions differ
RepeatConditional	conditioned stack — Dreamer 4’s transformer
Tokenwise	apply a module per token
InitWith	override the initializer for a subtree
StopGradParams	freeze a subtree’s parameters

Tokenwise and eval-mode state

Setting an attribute through Tokenwise did not propagate to the wrapped module for a while, which meant BatchNorm stayed in training mode during planning. If you wrap a module carrying eval-mode state, verify the mode actually reaches it.

Pre-built models

The models are comptime compositions of the primitives and combinators above, exported from noeira.nn.models:

from noeira.nn.models import ResNet18Backbone, GPT, ViT, PatchEmbed
from noeira.nn.models import ResBlockConv2DBN, ResBlockDownsampleBN

Model	Notes
ResBlockConv2DBN · ResBlockDownsampleBN	CIFAR ResNet blocks — ResNet-20 on CIFAR-10 in the examples
ResNet18Backbone[IN_CH, H, W]	ImageNet ResNet-18 topology, head-less — the ACT image backbone
GPT · GPTDrop · GPTDropTied	decoder-only transformer; the tied variant shares the embedding
ViT · PatchEmbed	vision transformer encoder with a classification head
TransformerBlock · MultiHeadAttention · TransformerFFN	the shared transformer pieces
Conv2DReLU · Conv2DBatchNormReLU · CifarFeatureClassifier	small convolutional trunks

ResNet18Backbone is torchvision’s resnet18 without its classifier: a 7×7 stride-2 stem and a 3×3 stride-2 max-pool, then four stages of two basic blocks, for a total stride of 32 and a [BATCH, 512 · H/32 · W/32] NCHW output (a 240×320 frame becomes 8×10 = 80 tokens). It uses ordinary trainable BatchNorm2D, not torchvision’s frozen batch norm. ImageNet weights load by name from timm/resnet18.tv_in1k through the map in noeira/nn/models/resnet18_torch.mojo, with no PyTorch install: the map renames torchvision’s parameters and zero-fills the convolution biases torchvision does not have.

ViT is PatchEmbed (a patch-stride Conv2D and a transpose) → a position bias → n_layers transformer blocks → LayerNorm → mean-pool → a linear head. SmolVLA’s SigLIP vision tower is a separate model in noeira/deep_agents/smolvla/vision.mojo, built from the same primitives.

pixi run -e nvidia mojo run -I . examples/nn/resnet/resnet20_cifar10_training_storage_gpu.mojo
pixi run -e nvidia mojo run -I . examples/nn/transformer/gpt_tinyshakespeare_training_gpu.mojo
pixi run -e nvidia mojo run -I . examples/nn/vit/vit_cifar_training_storage_gpu.mojo

ComputeGraph

Sequential covers chains. When a loss has a genuine DAG shape — several inputs, shared subgraphs, multiple outputs — use the named-node builder in combinators/compute_graph.mojo.

from noeira.nn.combinators.compute_graph import ComputeGraph

Nodes are declared with names and wired by reference, and the graph resolves the backward order itself. This is what the world-model losses use, where a single update touches an encoder, a dynamics net, several prediction heads and a KL term that has to flow to two places.

For a plain feed-forward network, Sequential is simpler and generates better code — reach for ComputeGraph when the topology actually needs it.

Compile-time cost

Deep parameterization is not free, and this is where it bites first.

Instantiation depth can explode compile times

Two known failure modes, both hit in this repository:

• A long chain of calls in one function (~80+) can trigger an inliner explosion. Extract helpers.
• A deep comptime loop building nested modules can hang the compiler outright. The LeWM depth-loop and eval-block cases both needed the body pulled into a separate function.

If a build stops making progress rather than failing, suspect instantiation depth before suspecting your code.

A related practical note: probe compilation with mojo build --emit=llvm rather than a full link. macOS ld rejects the very long generic symbol names this design produces, so a link failure can be a false negative that has nothing to do with your code.

Next

• Primitives — the leaves to compose
• Optimizers and losses — training the result
• Compile-time composition — the idiom behind all of this