Skip to content
noeiranoeiranoeira

Neural networks

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

noeira.nn is the deep-learning framework the RL agents are built on — described in the project’s own words as “something between Tinygrad and PyTorch”. It is usable on its own: the MNIST, CIFAR-10, ResNet, GPT and ViT examples never touch the RL code.

Networks are assembled from small reusable Module blocks, each Param carrying its own value and gradient tensors. Composition happens at compile time, so a network is a type — see compile-time composition.

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


var trainer = Trainer[Net, NC, IN_DIM, BATCH, "gpu"].make[Kaiming](
    Optional(ctx), lr=1e-3
)

What is in it

Section titled “What is in it”

Modules & Params

The storage design — a Param owns four tensors, and reflection over them is what makes optimizers and checkpointing work without registration boilerplate.

70+ primitives

The layer catalogue — Linear, Conv2D, self- and cross-attention, RoPE, LSTM/GRU, normalizations, and RL-specific ops.

Combinators & models

Composition — Sequential, Residual, Parallel, Repeat, plus ResNet-18, GPT and ViT, and the ComputeGraph DAG builder.

Optimizers & losses

SGD, Adam, AdamW, gradient clipping, LR schedules, and the loss and initializer catalogues.

Training

The supervised Trainer — whole-dataset GPU runs, AMP, optional CUDA-graph capture, and the bundled datasets.

Distributed training

DDP and ZeRO-1/2/3 behind one interface, over N GPUs from one process — or every rank simulated on one device.

The design in one page

Section titled “The design in one page”

Modules own their weights. A Module holds Param fields directly. There is no external parameter registry and no parameters() call to remember — a default for_each_param implementation reflects over the struct’s Param-typed fields, which is how the optimizer and the checkpoint layer find them.

A Param is four tensors. val and grd always; m and v — Adam’s moments — allocated lazily on the first stateful optimizer step. Co-locating moment state with the parameter rather than in a side table means checkpointing rides the same walk and exact training resume comes free.

Every tensor carries CPU and GPU storage. Tensor holds both a host List and a device buffer; the target parameter decides which is live. This is why the same module source compiles for both.

Forward and backward are explicit. Each module implements forward and vjp (the vector-Jacobian product). There is no tape: the backward pass is the composition of the modules’ vjps, resolved at compile time.

Fusion is a type. Adjacent operations that could fuse have a fused counterpart, and choosing it is how you opt in:

Linear[D, H], ReLU[H]   # two modules, two kernels
LinearReLU[D, H]        # one module, one fused kernel

MatMul+Bias+Activation collapses into a single GPU kernel this way (LinearAct and its aliases). Nothing fuses automatically across modules, and convolutions have no fused-activation form. Inside a module, the NVIDIA path folds the bias, and where it measured faster the activation, into the GEMM’s epilogue; see where the GEMMs run.

A complete example

Section titled “A complete example”

MNIST to 97%+ in five epochs, on GPU (examples/nn/mlp/mlp_mnist_training_storage_gpu.mojo):

from max.gpu.host import DeviceContext
from noeira.nn.datasets import MNIST
from noeira.nn.constants import DT
from noeira.nn.core.initializer import Kaiming
from noeira.nn.primitives.linear import Linear
from noeira.nn.primitives.linear_relu import LinearReLU
from noeira.nn.combinators.sequential import Sequential
from noeira.nn.training.trainer import Trainer




def main() raises:
    comptime IN_DIM = 784
    comptime NC = 10
    comptime BATCH = 100


    var ds = MNIST()
    var c = DeviceContext()


    comptime Net = Sequential[
        LinearReLU[IN_DIM, 256],
        LinearReLU[256, 128],
        Linear[128, NC],
    ]


    var trainer = Trainer[Net, NC, IN_DIM, BATCH, "gpu"].make[Kaiming](
        Optional(c), lr=1e-3
    )


    # one-hot the labels
    var train_y = List[Scalar[DT]](length=MNIST.N_TRAIN * NC, fill=0.0)
    for i in range(MNIST.N_TRAIN):
        train_y[i * NC + Int(ds.train_labels[i])] = 1.0


    var result = trainer.train_gpu[MNIST.N_TRAIN, MNIST.N_TEST](
        ds.train_images, train_y, ds.test_images, ds.test_labels,
        Optional(c), epochs=5, shuffle=True,
    )
Terminal window
pixi run -e apple  mojo run -I . examples/nn/mlp/mlp_mnist_training_storage_gpu.mojo
pixi run -e nvidia mojo run -I . examples/nn/mlp/mlp_mnist_training_storage_gpu.mojo

The dataset uploads to the device once and each batch is a zero-copy sub-buffer view — no per-batch host-to-device transfer.

Package layout

Section titled “Package layout”
noeira/nn/
├── core/          Module trait, Param, Tensor, initializers, checkpoint, AMP,
│                  safetensors + PyTorch name maps, split-K GEMM,
│                  cuBLAS / cuBLASLt / cuDNN bindings
├── primitives/    70+ leaf layers, fused attention kernels
├── combinators/   Sequential, Residual, Parallel, Repeat, Checkpointed, ComputeGraph
├── models/        ResNet blocks, ResNet-18, GPT, ViT, conv and CIFAR feature nets
├── optimizer/     SGD, Adam, AdamW, grad clipping, LR schedules
├── loss/          CrossEntropy, MSE, SoftCE, TwoHot, GaussianNLL
├── training/      Trainer, autoregressive trainer, augmenter
├── distributed/   process group, DDP, ZeRO-1/2/3, the distributed trainer
├── datasets/      MNIST, CIFAR-10, TinyShakespeare, lewm_pusht
└── random/        host RNG (box-muller, gaussian)

Next

Section titled “Next”