EXAMPLES

Building Blocks

Tier 2: reusable neural-network building blocks expressed as config factory functions. Each file provides a plain Python function that returns a laco config node, not a live module, so it can be composed into larger config trees without instantiation.

All source files live under sources/laco/examples/blocks/.

Design philosophy

Building-block files follow a strict convention:

Define a plain nn.Module class: the implementation, written without any laco awareness.
Expose a config factory function: a short function (e.g. basic_block(...)) that calls L.call(TheModule)(...) with the given arguments and returns the resulting config node.
Export both via __all__: callers can import either the class (for direct use in custom code) or the factory (for laco composition).

The factory function is the laco entry point. Calling it is pure Python; it runs at config-build time and produces a config node that will instantiate TheModule when laco.instantiate is called.

1. `blocks/residual.py`: ResNet-style residual blocks

Source: sources/laco/examples/blocks/residual.py

Provides two block variants (BasicBlock, Bottleneck) and their corresponding config factories (basic_block, bottleneck_block). Used by laco.examples.models.resnet.

Key source patterns

import laco.language as L
from torch import nn

__all__ = ["basic_block", "bottleneck_block", "BasicBlock", "Bottleneck"]


class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(self, in_channels: int, channels: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv2d(
            in_channels, channels, 3, stride=stride, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)
        self.relu = nn.ReLU(inplace=True)
        if stride != 1 or in_channels != channels:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, channels, 1, stride=stride, bias=False),
                nn.BatchNorm2d(channels),
            )
        else:
            self.downsample = nn.Identity()

    def forward(self, x):
        identity = self.downsample(x)
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return self.relu(out + identity)


class Bottleneck(nn.Module):
    expansion: int = 4

    def __init__(self, in_channels: int, channels: int, stride: int = 1):
        super().__init__()
        out_channels = channels * self.expansion
        self.conv1 = nn.Conv2d(in_channels, channels, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(
            channels, channels, 3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(channels)
        self.conv3 = nn.Conv2d(channels, out_channels, 1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        if stride != 1 or in_channels != out_channels:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels),
            )
        else:
            self.downsample = nn.Identity()

    def forward(self, x):
        identity = self.downsample(x)
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        return self.relu(out + identity)


# Config factory functions: these are what downstream configs import
def basic_block(in_channels: int, channels: int, stride: int = 1):
    return L.call(BasicBlock)(
        in_channels=in_channels, channels=channels, stride=stride
    )


def bottleneck_block(in_channels: int, channels: int, stride: int = 1):
    return L.call(Bottleneck)(
        in_channels=in_channels, channels=channels, stride=stride
    )

Annotated walkthrough

BasicBlock and Bottleneck are ordinary nn.Module subclasses. They contain no laco imports; this separation is intentional. The module implementation is testable without a config system.

Conditional downsample in __init__nn.Identity() vs. nn.Sequential(Conv2d, BN) is chosen at construction time based on stride and channel dimensions. This is standard PyTorch: laco does not interfere with the __init__ body.

basic_block(in_channels, channels, stride): the factory A single-line wrapper: return L.call(BasicBlock)(in_channels=in_channels, ...). The return value is a config node, not a BasicBlock instance. Callers can embed it in a larger config tree, pass it to L.repeat, or override its fields before instantiation.

bottleneck_block follows the same pattern with expansion=4 (four times as many output channels as channels).

Composing into a larger config

from laco.examples.blocks.residual import basic_block
import laco.language as L
from torch import nn

# Stack four identical residual blocks
backbone = L.call(nn.Sequential, expand_args=True)(
    L.repeat(4, basic_block(in_channels=64, channels=64))
)

# Mix strides for a ResNet-style stage
stage = L.call(nn.Sequential)(
    basic_block(in_channels=64, channels=128, stride=2),  # downsample
    basic_block(in_channels=128, channels=128),            # maintain
    basic_block(in_channels=128, channels=128),
)

What this demonstrates

Config factory over bare nn.Module: L.call(BasicBlock)(...) returns a config node
expansion class attribute propagates through __init__ without laco involvement
Conditional construction (downsample path) is pure Python, not a laco concern
Both factory and class are exported; downstream code chooses the right entry point

2. `blocks/transformer.py`: Pre-norm Transformer encoder block

Source: sources/laco/examples/blocks/transformer.py

Vision-style pre-norm Transformer encoder block: LayerNorm → MultiheadAttention → residual → LayerNorm → MLP → residual. Used by ViT-style image encoders.

Key source patterns

import laco.language as L
from torch import nn

__all__ = ["transformer_block", "TransformerBlock"]


class TransformerBlock(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size)
        self.attn = nn.MultiheadAttention(
            hidden_size, num_heads, dropout=dropout, batch_first=True
        )
        self.norm2 = nn.LayerNorm(hidden_size)
        intermediate = int(hidden_size * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, intermediate),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate, hidden_size),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        h = self.norm1(x)
        attn_out, _ = self.attn(h, h, h, need_weights=False)
        x = x + attn_out
        return x + self.mlp(self.norm2(x))


def transformer_block(
    hidden_size: int,
    num_heads: int,
    mlp_ratio: float = 4.0,
    dropout: float = 0.0,
):
    return L.call(TransformerBlock)(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        dropout=dropout,
    )

Annotated walkthrough

Pre-norm structureLayerNorm is applied before attention and MLP (pre-norm), not after (post-norm). This is the ViT/DINOv3 convention. The forward method is standard PyTorch; laco is not involved.

mlp_ratio: derived dimensionintermediate = int(hidden_size * mlp_ratio) is computed in __init__, not stored in the config. The config stores only hidden_size and mlp_ratio; the derived integer is an implementation detail.

Full self-attention with nn.MultiheadAttention Uses standard nn.MultiheadAttention with batch_first=True (sequences first dimension). No RoPE or GQA; for LLM-style attention see blocks/decoder.py.

transformer_block(...): the factory Four arguments map directly to TransformerBlock.__init__. The factory has default values matching the class, so callers only need to specify the required arguments.

Composing into a larger config

from laco.examples.blocks.transformer import transformer_block
import laco.language as L
from torch import nn

# 12-layer ViT encoder
encoder = L.call(nn.Sequential, expand_args=True)(
    L.repeat(
        12,
        transformer_block(hidden_size=768, num_heads=12, mlp_ratio=4.0, dropout=0.1),
    )
)

What this demonstrates

Attention + FFN composition expressed as a single config factory
Derived dimensions (intermediate) computed in __init__, not in the config
Pre-norm convention with nn.LayerNorm
L.repeat for stacked encoder blocks in downstream configs

3. `blocks/decoder.py`: Autoregressive LLM decoder block

Source: sources/laco/examples/blocks/decoder.py

LLM-style pre-norm decoder block: RMSNorm → GroupedQueryAttention (causal, RoPE) → residual → RMSNorm → SwiGLU → residual. Used by laco.examples.models.qwen3 and laco.examples.models.gemma3.

Key source patterns

import laco.language as L
from laco.examples.layers.gqa_attention import GroupedQueryAttention
from laco.examples.layers.rms_norm import RMSNorm
from laco.examples.layers.swiglu import SwiGLU
from torch import nn

__all__ = ["decoder_block", "DecoderBlock"]


class DecoderBlock(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        intermediate_size: int,
        rms_norm_eps: float = 1e-6,
        rope_base: float = 10_000.0,
        sliding_window: int | None = None,
    ):
        super().__init__()
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.self_attn = GroupedQueryAttention(
            hidden_size=hidden_size,
            num_heads=num_heads,
            num_kv_heads=num_kv_heads,
            causal=True,
            use_rope=True,
            rope_base=rope_base,
            sliding_window=sliding_window,
        )
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.mlp = SwiGLU(hidden_size, intermediate_size)

    def forward(self, x):
        x = x + self.self_attn(self.input_layernorm(x))
        return x + self.mlp(self.post_attention_layernorm(x))


def decoder_block(
    hidden_size: int,
    num_heads: int,
    num_kv_heads: int,
    intermediate_size: int,
    rms_norm_eps: float = 1e-6,
    rope_base: float = 10_000.0,
    sliding_window: int | None = None,
):
    return L.call(DecoderBlock)(
        hidden_size=hidden_size,
        num_heads=num_heads,
        num_kv_heads=num_kv_heads,
        intermediate_size=intermediate_size,
        rms_norm_eps=rms_norm_eps,
        rope_base=rope_base,
        sliding_window=sliding_window,
    )

Annotated walkthrough

GQA (Grouped-Query Attention)num_kv_heads < num_heads enables grouped-query attention (GQA). This reduces the KV-cache size proportionally. The ratio is stored explicitly in the config, so model-specific variants (Qwen3, Gemma3) can override it independently.

sliding_window: int | NoneNone → full causal attention (used by Qwen3 and Gemma3 odd-indexed layers). An integer value enables banded causal attention with that window size (used by Gemma3 even-indexed layers). Because L.call stores the value verbatim, a config override can switch this per-layer.

RMSNorm and SwiGLU Custom layer implementations from laco.examples.layers. Like the DecoderBlock itself, these are plain nn.Module subclasses; no laco awareness is required inside the implementation.

Composing into a larger config

from laco.examples.blocks.decoder import decoder_block
import laco.language as L
from torch import nn

# Uniform full-attention stack (Qwen3-style)
decoder_stack = L.call(nn.Sequential, expand_args=True)(
    L.repeat(
        28,
        decoder_block(
            hidden_size=1024,
            num_heads=16,
            num_kv_heads=8,
            intermediate_size=3072,
        ),
    )
)

# Alternating sliding-window / full-attention (Gemma3-style)
# constructed with explicit Python iteration at config-build time
import laco.language as L

layers = []
for i in range(28):
    window = None if i % 2 == 1 else 4096
    layers.append(
        decoder_block(
            hidden_size=1024,
            num_heads=16,
            num_kv_heads=4,
            intermediate_size=4096,
            sliding_window=window,
        )
    )
decoder_stack = L.call(nn.Sequential, expand_args=True)(*layers)

What this demonstrates

GQA config with separate num_heads / num_kv_heads fields
sliding_window: int | None: optional field stored verbatim, overridable per-layer
Custom layer types (RMSNorm, SwiGLU, GroupedQueryAttention) as transparent config targets
Python-level iteration for heterogeneous stacks (alternating attention patterns)

Design patterns for building blocks

1. Factories return config nodes, not modules

# Correct: returns a config node
def basic_block(in_channels, channels, stride=1):
    return L.call(BasicBlock)(in_channels=in_channels, channels=channels, stride=stride)

# Wrong: instantiates immediately, bypasses config system
def basic_block(in_channels, channels, stride=1):
    return BasicBlock(in_channels, channels, stride)  # do not do this

The factory is called at config-build time (import time for the consuming config). Instantiation happens only when laco.instantiate is called on the root config.

2. `L.repeat` for homogeneous stacks

# N identical blocks, all with the same hyperparameters
stack = L.call(nn.Sequential, expand_args=True)(
    L.repeat(num_layers, transformer_block(hidden_size=768, num_heads=12))
)

expand_args=True is required when passing a L.repeat list to nn.Sequential, which takes *modules not a list.

3. Python iteration for heterogeneous stacks

When layers differ (different strides, alternating window sizes), use Python for/if at config-build time:

layers = []
for i in range(num_stages):
    stride = 2 if i == 0 else 1
    layers.append(basic_block(in_channels=ch, channels=ch, stride=stride))
stack = L.call(nn.Sequential, expand_args=True)(*layers)

4. Export both class and factory

__all__ = ["transformer_block", "TransformerBlock"]

Downstream code that does not use laco can import TransformerBlock directly. Laco configs import the factory. Both are tested independently.

Foundation Examples

Typed-Group Variants