netcl.nn: Modules, Layers, ResNet

netcl.nn is the model-building layer of netcl. It is shaped after PyTorch's torch.nn but tuned for the OpenCL pipeline: every Module is a Tensor container, every Parameter is just a Tensor with requires_grad=True, and the heavy lifting happens in the JIT Compiler through the ops API and the autograd API.

All public symbols are re-exported from the package root:

from netcl.nn import (
    Module, Linear, Conv2d, BatchNorm2d, ReLU, LeakyReLU, Sigmoid, Tanh,
    Dropout, MaxPool2d, Flatten, LayerNorm, MultiheadAttention,
    TransformerEncoderLayer, Sequential, Parameter,
    MSELoss, L1Loss, SmoothL1Loss, BCELoss, BCEWithLogitsLoss,
    mse_loss, cross_entropy,
    xavier_uniform, kaiming_uniform, constant,
    build_sequential, example_mlp_config, example_cnn_config,
)
# Decoder-LM and Modern Transformer modules:
from netcl.nn.transformer import CausalSelfAttention, TransformerBlock, PositionalEmbedding
from netcl.nn.modern import RMSNorm, RotaryAttention, SwiGLU, DecoderBlock, KVCache, TiedLinear
# ResNet18 is lazy-loaded to avoid pulling in the full model at import time:
from netcl.nn import ResNet18

Class Hierarchy

All neural network layers and architectures inherit from nn.Module.

nn.Module

Module is the base class. You subclass it, assign submodules and parameters in __init__, and implement forward(x). The training loop just calls model(x), which is implemented as self.forward(x).

from netcl.nn import Module, Linear
import netcl.autograd as ag

class MyNet(Module):
    def __init__(self, queue):
        super().__init__()
        self.fc1 = Linear(queue, 784, 256)
        self.fc2 = Linear(queue, 256, 10)

    def forward(self, x):
        x = ag.relu(self.fc1(x))
        return self.fc2(x)

nn.Parameter

Parameter is a thin wrapper around a Tensor that is automatically registered with the parent Module when assigned as an attribute. The Tensor it wraps has requires_grad=True by default.

from netcl.nn import Parameter
from netcl.core.device import manager

q = manager.default("auto").queue
w = Parameter.from_shape(q, (10, 784), dtype="float32")
# Once assigned to a Module attribute, `w` appears in module.parameters().

In practice you almost never instantiate Parameter by hand, layer constructors (Linear, Conv2d, BatchNorm etc.) build the right parameters for you with the right initial values from init.

nn.Linear

from netcl.nn import Linear
layer = Linear(queue, in_features, out_features, bias=True)
# queue is optional: omit it and Linear auto-discovers the default device
layer = Linear(in_features=784, out_features=256)

A fully-connected layer that computes y = x @ W^T + b. W has shape (out_features, in_features); b has shape (out_features,). Both are Parameters, initialized with Kaiming normal for W and zeros for b. The forward pass dispatches to the matmul op (auto-tuned via the KernelSelector and the JIT Compiler) plus a bias_add.

nn.Conv2d

from netcl.nn import Conv2d
layer = Conv2d(
    queue,           # optional: omit to auto-discover the default device
    in_channels, out_channels, kernel_size,
    stride=1, pad=0, bias=True,
)
# Example: 3→16 channels, 3×3 kernel, same-padding
layer = Conv2d(queue, 3, 16, kernel_size=3, stride=1, pad=1)

Note: the padding parameter is pad=, not padding=. There is no dilation or groups support.

2D convolution over NCHW tensors. Implemented as im2col plus a matmul by default; the optimized variant in ops/conv2d_optimized.py is selected by the KernelSelector when the workload and device warrant it. The 1×1 specialization goes through implicit-GEMM, and 3×3 stride-1 maps to the Winograd fused variant unless NETCL_CONV_WINOGRAD=0.

nn.MaxPool2d

2D max-pooling layer. Has a backward kernel registered in autograd/ops.py so it is fully differentiable.

from netcl.nn import MaxPool2d
pool = MaxPool2d(kernel_size=2, stride=2)

nn.BatchNorm2d

Per-channel normalization. In training mode it computes batch statistics; in eval mode it uses the running mean and variance. The fused variant combines BN+ReLU into a single kernel (see fused-ops table).

from netcl.nn import BatchNorm2d
bn = BatchNorm2d(queue, num_features=64, eps=1e-5, momentum=0.1)
# queue must belong to the same OpenCL context as the input tensor

When model.eval() is set, bn uses its running_mean / running_var buffers and skips the batch-statistics update.

nn.LayerNorm

Normalizes across the last len(normalized_shape) dimensions of the input. No running-statistics state.

from netcl.nn import LayerNorm
ln = LayerNorm(normalized_shape=(64,))

nn.Dropout

Stochastic regularization. Inactive when model.eval() is in effect.

from netcl.nn import Dropout
drop = Dropout(p=0.5)

Transformer Layers (nn.MultiheadAttention & nn.layers)

nn.MultiheadAttention

Multi-head scaled dot-product attention over queries, keys, and values:

from netcl.nn import MultiheadAttention

mha = MultiheadAttention(embed_dim=256, num_heads=8, queue=q)
out = mha(query, key, value)

Projections for Q, K, V and output projection are performed via OpenCL Linear layers, followed by batched matrix multiplication (bmm) and softmax.

nn.TransformerEncoderLayer

A standard post-LayerNorm encoder block: 1. src + MultiheadAttention(src, src, src) -> LayerNorm 2. x + Linear2(ReLU(Linear1(x))) -> LayerNorm

from netcl.nn import TransformerEncoderLayer

layer = TransformerEncoderLayer(d_model=256, nhead=8, dim_feedforward=1024, queue=q)
out = layer(src)

Decoder-LM Blocks (netcl.nn.transformer)

Decoder-oriented blocks featuring causal masking and pre-LN residual connections:

CausalSelfAttention

Multi-head self-attention with a cached upper-triangular causal mask:

from netcl.nn.transformer import CausalSelfAttention

attn = CausalSelfAttention(embed_dim=256, num_heads=8, causal=True, dropout=0.1, queue=q)
out = attn(x)

Uses a fused 3 * embed_dim projection for Q, K, V to reduce GEMM kernel launches 3x.

TransformerBlock

Pre-LN decoder block: * h = x + CausalSelfAttention(LayerNorm(x)) * out = h + MLP(LayerNorm(h)) (supporting GELU, ReLU, or SiLU activations).

from netcl.nn.transformer import TransformerBlock

block = TransformerBlock(d_model=256, num_heads=8, dim_feedforward=1024, activation="gelu", queue=q)
out = block(x)

PositionalEmbedding

Learned absolute positional embeddings:

from netcl.nn.transformer import PositionalEmbedding

pos_emb = PositionalEmbedding(max_len=2048, d_model=256, queue=q)
emb = pos_emb(seq_len=128)

Modern Decoder-LM Stack (netcl.nn.modern)

The modern decoder stack powering architectures like LLaMA and Mistral:

RMSNorm

Root Mean Square Layer Normalization: y = x * gamma / sqrt(mean(x^2) + eps) over the last axis.

One memory pass over rows instead of two, with no additive bias (beta), cutting memory bandwidth in half compared to LayerNorm:

from netcl.nn.modern import RMSNorm

norm = RMSNorm(dim=512, eps=1e-6, queue=q)
out = norm(x)

RotaryAttention

Causal self-attention combining: 1. RoPE (Rotary Position Embeddings): Query and key rotation past context length without learned parameters. 2. QK-Norm: RMSNorm applied to query and key heads before attention to prevent logit explosion. 3. Streaming FlashAttention: Executes attention without materializing the (B, H, S, S) matrix. 4. KVCache Integration: Transition between multi-token prompt prefill and single-token decode.

from netcl.nn.modern import RotaryAttention

attn = RotaryAttention(dim=512, n_head=8, rope_base=10000.0, qk_norm=True, queue=q)
out = attn(x, cache=kv_cache, layer=0)

KVCache

Pre-allocated key/value buffer per layer for O(1) token insertion and O(n) autoregressive decoding:

from netcl.nn.modern import KVCache

cache = KVCache(n_layer=12, batch=1, n_head=8, head_dim=64, capacity=2048, queue=q)
cache.reset()

SwiGLU

Swish Gated Linear Unit: down(silu(gate(x)) * up(x)).

from netcl.nn.modern import SwiGLU

ffn = SwiGLU(dim=512, hidden=1376, queue=q)  # default hidden = 8/3 * dim
out = ffn(x)

DecoderBlock

Complete pre-norm modern transformer block integrating RMSNorm, RotaryAttention, and SwiGLU:

from netcl.nn.modern import DecoderBlock

block = DecoderBlock(dim=512, n_head=8, rope_base=10000.0, qk_norm=True, queue=q)
out = block(x, cache=kv_cache, layer=0)

TiedLinear

Ties the language model's output projection head directly to its token embedding matrix:

from netcl.nn.modern import TiedLinear

lm_head = TiedLinear(embedding_weight)
logits = lm_head(hidden_states)

nn.build_sequential

Builds a Sequential model from a list of layer config dicts. Useful for config-file-driven architectures or hyperparameter search over layer widths.

from netcl.nn import build_sequential, example_mlp_config
from netcl.core.device import manager

q = manager.default("auto").queue

# Use a built-in config template
config = example_mlp_config(input_dim=784, hidden=256, num_classes=10)
model = build_sequential(q, config)

# Or write your own config
config = [
    {"type": "Linear", "args": {"in_features": 784, "out_features": 256}},
    {"type": "ReLU",   "args": {}},
    {"type": "Linear", "args": {"in_features": 256, "out_features": 10}},
]
model = build_sequential(q, config)

build_sequential_from_json(queue, path) does the same from a JSON file.

nn.ResNet18

Pre-built ResNet-18 model. Forward expects NCHW input. Default num_classes=10 (CIFAR-friendly); override for ImageNet heads.

from netcl.nn import ResNet18
from netcl.core.device import manager

q = manager.default("auto").queue
model = ResNet18(queue=q, num_classes=10)

Internally built from Conv2d, BatchNorm2d, ReLU, and the residual add, all dispatched to fused kernels where supported (see fused-ops table).

Fused Ops

The nn.functional wrappers call into autograd/ops.py, so every fused op is fully differentiable and participates in the Tape.

Initialization (nn.init)

Initializers are in netcl.nn.init and are also re-exported from netcl.nn. They work in-place on a Tensor parameter.

from netcl.nn import xavier_uniform, kaiming_uniform, constant

xavier_uniform(model.fc1.weight)      # Tanh-family activations
kaiming_uniform(model.fc1.weight)     # ReLU-family (the default for Linear/Conv2d)
constant(model.bn.weight, 1.0)        # any constant fill

Functional (nn.functional)

Stateless variants of activations and losses. These wrap the same autograd ops that the layer classes call, so they are fully differentiable under a Tape.

All functional ops wrap the autograd engine and return Node objects, not plain Tensors. Access the underlying value via .value:

from netcl.nn import functional as F

y    = F.relu(x)             # returns Node; use y.value.to_host() to read back
y    = F.sigmoid(x)
loss = F.cross_entropy(logits, targets)   # targets: integer Tensor (int32)
loss = F.mse_loss(pred, target)
loss = F.binary_cross_entropy(pred, target)

Putting It Together: a Small Classifier

import numpy as np
import netcl.autograd as ag
from netcl.core.device import manager
from netcl.nn import Linear, ReLU, Sequential, Dropout
from netcl.nn import functional as F
from netcl.optim import Adam

q = manager.default("auto").queue
model = Sequential(
    Linear(q, 784, 256), ReLU(), Dropout(p=0.1),
    Linear(q, 256, 128), ReLU(),
    Linear(q, 128, 10),
)
opt = Adam(model.parameters(), lr=3e-4)

for x, y in loader:                       # x: Tensor (B, 784),  y: Tensor (B,) int32
    with ag.Tape() as tape:
        logits = model(ag.tensor(x))       # wrap input so autograd tracks it
        loss = F.cross_entropy(logits, y)  # loss is a Node
    tape.backward(loss)
    opt.step()
    opt.zero_grad()
    print(f"loss = {loss.value.to_host()[0]:.4f}")

See also

• Tensor: the Tensor that Module wraps.
• ops API: the elementwise, matmul, and conv2d primitives that the layers call into.
• autograd API: the Tape and Node that power the backward pass.
• JIT Compiler: how the elementwise, matmul, and conv kernels are compiled and fused.
• MNIST with MLP: a complete training example that uses Linear, Dropout, and Sequential.
• core API: the DeviceManager and Tensor foundation.
• optim API: the SGD / Adam / AdamW optimizers.