Skip to content

Adding an operator

Every HNDL operation is one nn.Module subclass with an @operator(...) declaration above it. The declaration is the single source of truth for the resolver (shape relation, arguments), the backend (the class itself), the tests (examples), and the documentation (summary, help text, docstring).

Files

An operator touches exactly these files; nothing shared is edited:

File Purpose
src/hndl/operators/<alias>.py The decorated class. Discovered automatically.
tests/operators/test_<alias>.py Operator-specific numerics and error cases.
docs/operators/<alias>.md Generated: run python -m hndl.docs. Do not edit by hand.

docs/operators/index.md is regenerated by the same command.

The declaration

from torch import nn

from ..operator import Arg, Example, operator


@operator(
    "layer_norm",                                   # alias used in configs and ops.<alias>
    summary="Normalize the last axis with a learned scale and bias.",
    shape="x[B, ..., D] -> out[B, ..., D]",         # ports and their relation
    args={
        "eps": Arg(float, 1e-5, min=0, exclusive_min=True, help="Added to the variance."),
        "affine": Arg(bool, True, positional=False, help="Learn per-feature scale and bias."),
    },
    examples=[
        Example("linear(64)\nlayer_norm()\nlinear()", ("B", 32), ("B", 10),
                "Normalizes the hidden features."),
    ],
    category="normalization",
)
class LayerNorm(nn.Module):
    """Markdown rendered into the docs: what it computes, axis conventions,
    train/eval behavior, parameter names."""

    def __init__(self, eps, affine, *, D):
        ...

    def forward(self, x):
        ...

Fields:

  • alias is what configs call. identity= defaults to the alias; version=1. Aliases become reserved names in configs, so avoid names users would want for local variables.
  • summary is one sentence for the index. The class docstring is the long description (Markdown).
  • shape declares ports and their relation in a small DSL:
  • name[dims] per port, inputs before ->, outputs after. The first input port receives the implicit current tensor.
  • B is the batch axis and must come first. Names such as C or D_in are symbols shared across the ports of one node: the same name means the same extent. 2*C is an exact integer multiple. Literals like 4 fix an axis. ... stands for the same run of middle axes on every port that uses it, so x[B, ...] -> out[B, ...] preserves any supported shape.
  • A port with no brackets, x -> first, rest, is unconstrained by the DSL; pair it with relation=.
  • x* declares a variadic input (x0, x1, …); declare an int argument named input_count.
  • out* declares a variadic output (out0, out1, …); it must be the sole output port and the declaration names the argument fixing the count with outputs_from="<argument>" — a sequence (ints/strs), whose length is the count, or a bounded int, whose value is. A count of zero leaves the single declared out port; any other count returns a tuple and clears current.
  • batch= says how the operator treats axis 0. The default, "shared", means batch passes through: every port of the node carries the same batch entry, and B in a pattern stands for that entry, which may be a multiple such as 2*B. An operator that moves tensors across the batch axis — concat(axis=0), chunk(dim=0) — declares batch="relation" and sets axis 0 on each port itself, with s.batch(port) and s.axis(port, 0, ...).
  • relation= is a function receiving a node view s for rules the DSL cannot express (convolution arithmetic, products, splits). It runs in every solver sweep and may only add facts: s.shape(port), s.rank(port, r), s.axis(port, i, value), s.equal(p, q), s.arg(name, value), s.product(p, q), s.interval(port, i, lo, hi), s.error(code, msg), s.batch(port), s.share_batch(*ports), s.args, s.inputs, and s.policy (the selected policy identity or None). Shared relations live in operators/_relations.py. Give shape_text= a one-line description for the docs.
  • args maps names to Arg(type, default, ...). Omit the default to require the value. inferable=True lets the resolver solve an omitted dimension; dim="D_out" ties it to a shape symbol in both directions. Types: int, float, bool, str, "pair" (int or two ints), "ints", "strs". Every Arg needs help. Positional order is the mapping order; mark secondary arguments positional=False.
  • examples are runnable configs with their input and output shapes. The harness resolves, builds, runs, and round-trips every example on every available device, and the docs render them with their shape tables.
  • Optional hooks: positional_rest="shape" (leftover positionals fill an ints argument), policies={"up2": Policy(identity, requires)}, validate=fn(args) for cross-field checks, finalize=fn(args, input_shapes, output_shapes) to store resolved values (must be idempotent), reference=fn(module) returning a handwritten equivalent the harness compares against.

The constructor receives every argument by keyword, plus any shape symbol it names as a keyword-only parameter (*, D) and, if requested, input_shapes / output_shapes. It is called under torch.device(device), so create tensors normally; parameters are then cast to the plan dtype. forward receives tensors in declared input-port order and returns one tensor, or a tuple/dict for multiple output ports.

Checklist

  1. Create src/hndl/operators/<alias>.py with the decorated class.
  2. Add tests/operators/test_<alias>.py: compare against the PyTorch reference on random shapes (cpu and cuda), cover backward inference through the operator, and assert the error codes for invalid arguments.
  3. Run PYTHONPATH=src python -m pytest -q from the checkout. In a git worktree the .venv editable install points at the main checkout, so PYTHONPATH=src is required.
  4. ruff check src tests examples.
  5. python -m hndl.docs and commit the generated page. python -m hndl.docs --check must pass.
  6. Commit, push, and open a PR titled Add <alias> operator.

Custom operators outside the package

The same decorator works on a registry:

from hndl import Arg, Registry
from hndl.torch import network

registry = Registry.builtins()

@registry.operator("my_silu", identity="example.silu", summary="SiLU activation.",
                   shape="x[B, ...] -> out[B, ...]")
class SiLU(nn.SiLU):
    pass

model = network("linear(64); my_silu(); linear()", input_shape=("B", 128), output_shape=("B", 10),
                registry=registry, device="cuda:0")

registry.add(cls) registers a class decorated elsewhere with hndl.operator.