Notebook

Typed Groups and Schemas

Prerequisites: 01.why-laco.ipynb through 05.loading-saving-cli.ipynb. Familiarity with Python dataclasses (@dataclass, fields()).

Dependencies: torch (for nn.Module and optim.Optimizer as nominal types).

@L.params describes configs as a namespace of hyperparameters with interpolation support. That works for scalars. It does not work for swappable components, where one field should be nn.ReLU or nn.GELU or nn.SiLU, validated at configuration time. This notebook covers the mechanism that handles that case.

By the end you will understand:

Why @L.params is insufficient for swappable components.
L.Group[T]: declaring a typed set of alternatives.
@L.config: a structured schema that is a real dataclass.
L.slot: linking a schema field to a group.
L.Defaults and L.bind: choosing a default entry.
L.chosen: referencing the selected entry in the config body.
The complete typed/mlp.py example from first principles.

import dataclasses

import laco
import laco.language as L
from omegaconf import OmegaConf
from torch import nn, optim

Section 1: The Limitation of `@L.params` for Swappable Components

Suppose the activation function in an MLP needs to be configurable. The naive approach with @L.params looks like this:

@L.params
class hps_naive:
    dim_in: int = 128
    dim_hidden: int = 256
    dim_out: int = 64
    # PROBLEM: we want `activation` to be a *swappable component*, but @L.params
    # is built for scalar hyperparameters. Storing a class object (nn.ReLU) is
    # rejected when the params block is materialised into a DictConfig: OmegaConf
    # only accepts primitive scalars, not arbitrary Python class references.
    activation: type[nn.Module] = nn.ReLU

# Statically, the attribute is just an interpolation string into the params block:
print("hps_naive.activation ->", repr(hps_naive.activation))

# Materialising the block surfaces the limitation: a class is not a valid value.
try:
    OmegaConf.to_yaml(hps_naive())   # hps_naive() builds the dict; nn.ReLU is rejected
except Exception as e:
    print(f"{type(e).__name__}: {e}".splitlines()[0])
    print("\n=> @L.params cannot hold a swappable component (a class/object).")
    print("   Section 2 fixes this with a typed L.Group instead.")

Output

hps_naive.activation -> '${hps_naive.activation}'
UnsupportedValueType: Value 'ReLU' is not a supported primitive type

=> @L.params cannot hold a swappable component (a class/object).
   Section 2 fixes this with a typed L.Group instead.

With @L.params there are three concrete problems:

Problem	Detail
No validated choice set	There is no enforcement that `activation=gelu` is a valid alternative; any string is accepted until runtime crashes
Runtime-only errors	A typo like `activation=relu_typo` fails at runtime, inside a training job, not at edit time
No static type checking on alternatives	The IDE cannot autocomplete `ActivationGroup.relu` because `ActivationGroup` does not exist

The @L.params mechanism is designed for scalar hyperparameters (learning rate, batch size, number of layers). For swappable object nodes, a choice among several pre-configured nn.Module variants, the Group API is what's needed.

Section 2: `L.Group[T]`, Declaring a Config Group

A config group is a named set of interchangeable config nodes that all instantiate to the same type T. Subclassing L.Group[T] declares such a group:

class ActivationGroup(L.Group[nn.Module]):
    """Swappable activation functions."""

    relu = L.call(nn.ReLU)()
    gelu = L.call(nn.GELU)()
    silu = L.call(nn.SiLU)()

What just happened at class-creation time?

Group.__init_subclass__ fired and walked every non-dunder class attribute.
Each attribute whose value is a DictConfig node (produced by L.call) was registered into Hydra's ConfigStore under group="activationgroup", name="relu" etc.
A GroupEntry(group, name, node) record was stored in the module-level _ENTRY_ORIGINS dict, keyed by the node's id(), so L.bind can later recover the group/name from the entry value alone.

The group name defaults to cls.__name__.lower(): ActivationGroup → "activationgroup".

# The class attribute IS the DictConfig node (lie-typed as nn.Module)
print("type(ActivationGroup.relu):", type(ActivationGroup.relu))

# Dump shows the recipe that will be instantiated
print("\n--- ActivationGroup.relu ---")
print(laco.dump(ActivationGroup.relu))

print("--- ActivationGroup.gelu ---")
print(laco.dump(ActivationGroup.gelu))

Output

type(ActivationGroup.relu): <class 'omegaconf.dictconfig.DictConfig'>

--- ActivationGroup.relu ---
{_convert_: all, _laco_: 1, _target_: torch.nn.ReLU}

--- ActivationGroup.gelu ---
{_convert_: all, _laco_: 1, _target_: torch.nn.GELU}

# Iterate over all registered entries
for entry in ActivationGroup.entries():
    print(f"  group={entry.group!r}  name={entry.name!r}")

Output

  group='activationgroup'  name='relu'
  group='activationgroup'  name='gelu'
  group='activationgroup'  name='silu'

Static safety

Because ActivationGroup.relu is a real class attribute (lie-typed as nn.Module), a typo like ActivationGroup.relu_typo raises AttributeError immediately in Python and is a static error in pyright, before you ever run the training script.

Compare this to @L.params where any string could silently reach the YAML override machinery.

# Typo is caught immediately at import time (or by pyright at edit time)
try:
    _ = ActivationGroup.relu_typo
except AttributeError as e:
    print(f"AttributeError: {e}")

Output

AttributeError: type object 'ActivationGroup' has no attribute 'relu_typo'

Section 3: `@L.config`, The Typed Schema

A schema describes the hyperparameter fields for a model. @L.config is a PEP 681 dataclass_transform decorator: it applies dataclasses.dataclass under the hood, so the result is a real dataclass with IDE-visible fields, but it also handles the special L.slot(...) field specifier.

@L.config
class MLPSchema:
    """Typed configuration schema for the MLP."""

    dim_in:     int       = 128
    dim_out:    int       = 128
    dim_hidden: int       = 256
    num_layers: int       = 3
    # L.slot links this field to ActivationGroup.
    # At runtime the default is rewritten to the string "${activation}"
    # (an OmegaConf interpolation that the defaults list fills in).
    # Statically, pyright sees type nn.Module — the lie-typing contract.
    activation: nn.Module = L.slot(ActivationGroup)

# @L.config applied dataclasses.dataclass — MLPSchema is a real dataclass
print("Is dataclass:", dataclasses.is_dataclass(MLPSchema))

print("\nFields:")
for f in dataclasses.fields(MLPSchema):
    print(f"  {f.name}: {f.type}  default={f.default!r}")

Output

Is dataclass: True

Fields:
  dim_in: <class 'int'>  default=128
  dim_out: <class 'int'>  default=128
  dim_hidden: <class 'int'>  default=256
  num_layers: <class 'int'>  default=3
  activation: <class 'torch.nn.modules.module.Module'>  default='${activation}'

How `L.slot` works

When @L.config is applied it calls _resolve_slots(target), which walks the class body looking for _SlotSpec instances. For each one it:

Infers the package name from the field name (e.g. field activation → package "activation").
Rewrites the class attribute to a _SlotRef, which is a str subclass holding "${activation}", an OmegaConf variable interpolation.
Stores the originating group name (activationgroup) so L.bind can validate that the bound entry actually belongs to the right group.

When Hydra composes a config that includes both the schema and a defaults-list entry for activationgroup, the interpolation ${activation} resolves to the composed entry node.

# After @L.config the field default is the interpolation string
raw_default = MLPSchema.__dataclass_fields__["activation"].default
print("activation default (runtime):", repr(raw_default))
# Pyright / IDE sees: nn.Module (the lie)
# Python holds: _SlotRef('${activation}', group='activationgroup')

Output

activation default (runtime): '${activation}'

Section 4: `L.Defaults`, The Defaults List

Hydra's defaults list tells the composition engine which entries to merge, and in what order. L.Defaults produces a validated Python object that Laco serializes into the defaults key recognized by Hydra.

defaults = L.Defaults(
    L.self_,                                            # include this file's own fields
    L.bind(MLPSchema.activation, ActivationGroup.relu), # activation slot → relu entry
)

print("defaults:", defaults)

Output

defaults: ['_self_', {'activationgroup@activation': 'relu'}]

Step-by-step breakdown of `L.bind`

L.bind(MLPSchema.activation, ActivationGroup.relu)
       ^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^
       slot (_SlotRef)       entry (DictConfig node)

At runtime L.bind does:

Resolve the entry origin: looks up id(ActivationGroup.relu) in _ENTRY_ORIGINS to recover GroupEntry(group='activationgroup', name='relu', node=...).
Validate the slot: checks that the slot's stored group name (MLPSchema.activation.group == 'activationgroup') matches the entry's group. A mismatch (e.g. L.bind(MLPSchema.activation, OptimGroup.adam)) raises TypeError immediately, at config-authoring time.
Produce a DefaultsBinding: DefaultsBinding(group='activationgroup', name='relu') which serializes to {"activationgroup": "relu"} in the Hydra defaults list.

The result: Hydra will compose the relu node into the package named "activation" (the field name), which the interpolation ${activation} then resolves to.

# What happens if we accidentally bind the wrong group?
class OptimizerGroup(L.Group[optim.Optimizer]):
    sgd  = L.partial(optim.SGD)(lr=1e-2, momentum=0.9)
    adam = L.partial(optim.Adam)(lr=1e-3)

try:
    bad = L.bind(MLPSchema.activation, OptimizerGroup.adam)
except TypeError as e:
    print(f"TypeError: {e}")

Output

TypeError: bind() type mismatch: slot is bound to group 'activationgroup' but entry 'adam' belongs to group 'optimizergroup'.

Section 5: `L.chosen`, Referencing the Selected Entry in the Config Body

Inside the model config body (the L.call(nn.Sequential)(...) tree), the code must reference whatever activation function was selected by the defaults list. Hard-coding L.call(nn.ReLU)() here would defeat the whole point of having a swappable group.

L.chosen(ActivationGroup) produces exactly this: an OmegaConf interpolation that resolves to the composed activation entry at instantiation time.

activation_ref = L.chosen(ActivationGroup)

# Runtime value: an interpolation string
print("runtime:", repr(activation_ref))
# Pyright sees: nn.Module (lie-typed as the group element type)

Output

runtime: '${activationgroup}'

The interpolation ${activation} matches the package name derived from the L.slot(ActivationGroup) field (whose field name is "activation").

When Hydra composes a config that has:

schema.activation = "${activation}" (from the slot)
defaults: [{activationgroup: relu}] (from L.bind)

then ${activation} resolves to the full relu DictConfig node, so every L.chosen call in the model body receives the correct activation.

Section 6: Full Walkthrough, `typed/mlp.py`

Now that each primitive is covered, walk through the complete file line by line. The source lives at sources/laco/examples/typed/mlp.py.

# ------------------------------------------------------------------
# Step 1 — Declare the group of swappable activations
# ------------------------------------------------------------------
class ActivationGroup(L.Group[nn.Module]):  # type: ignore[no-redef]  # redefine for clarity
    """Swappable activation functions."""
    relu = L.call(nn.ReLU)()
    gelu = L.call(nn.GELU)()
    silu = L.call(nn.SiLU)()

# Each attribute is a DictConfig node registered in the ConfigStore.
# Group name: 'activationgroup'  (class name lower-cased)

# ------------------------------------------------------------------
# Step 2 — Declare the typed schema
# ------------------------------------------------------------------
@L.config
class MLPSchema:  # type: ignore[no-redef]
    """Typed schema for the MLP example."""
    dim_in:     int       = 128
    dim_out:    int       = 128
    dim_hidden: int       = 256
    num_layers: int       = 3
    activation: nn.Module = L.slot(ActivationGroup)
    # After @L.config, the activation field default is '${activation}'.
    # Pyright sees: nn.Module  (lie — matches the Group element type)

# ------------------------------------------------------------------
# Step 3 — Expose the schema as a module-level name
# ------------------------------------------------------------------
# When laco.load() reads this file, 'schema' is exposed as a fragment:
#   laco.load('configs://examples/typed/mlp.py#schema')
schema = MLPSchema

# ------------------------------------------------------------------
# Step 4 — Declare the defaults list
# ------------------------------------------------------------------
defaults = L.Defaults(
    L.self_,                                            # include this config's own fields
    L.bind(MLPSchema.activation, ActivationGroup.relu), # default activation = relu
)
# Serialises to: ["_self_", {"activationgroup": "relu"}]
print("defaults:", defaults)

Output

defaults: ['_self_', {'activationgroup@activation': 'relu'}]

# ------------------------------------------------------------------
# Step 5 — Build the model config using L.chosen and L.ref
# ------------------------------------------------------------------
model = L.call(nn.Sequential, root=True)(
    L.OrderedDict(
        (
            "input",
            L.call(nn.Sequential)(
                L.call(nn.Linear)(
                    in_features=L.ref("${schema.dim_in}"),
                    out_features=L.ref("${schema.dim_hidden}"),
                ),
                L.chosen(ActivationGroup),   # <-- resolves to selected activation
            ),
        ),
        (
            "hidden",
            L.call(nn.Sequential, expand_args=True)(
                L.repeat(
                    L.ref("${schema.num_layers}"),
                    L.call(nn.Sequential)(
                        L.call(nn.Linear)(
                            in_features=L.ref("${schema.dim_hidden}"),
                            out_features=L.ref("${schema.dim_hidden}"),
                        ),
                        L.chosen(ActivationGroup),  # same interpolation — stays in sync
                    ),
                ),
            ),
        ),
        (
            "output",
            L.call(nn.Sequential)(
                L.call(nn.Linear)(
                    in_features=L.ref("${schema.dim_hidden}"),
                    out_features=L.ref("${schema.dim_out}"),
                )
            ),
        ),
    )
)

print("model type (runtime):", type(model))

Output

model type (runtime): <class 'omegaconf.dictconfig.DictConfig'>

# The example at sources/laco/examples/typed/mlp.py declares the group, schema,
# defaults list, and model body. To load + compose it here we materialise the
# same wiring to a small file (the defaults list selects the activation entry,
# and L.chosen / L.slot share the "activation" package so the interpolation
# resolves). This mirrors the source file with the group binding made explicit.
import tempfile, textwrap, pathlib

_MLP_TYPED = textwrap.dedent("""
    import laco.language as L
    from torch import nn

    class ActivationGroup(L.Group[nn.Module]):
        relu = L.call(nn.ReLU)()
        gelu = L.call(nn.GELU)()
        silu = L.call(nn.SiLU)()

    @L.config
    class MLPSchema:
        dim_in:     int       = 128
        dim_out:    int       = 128
        dim_hidden: int       = 256
        num_layers: int       = 3
        activation: nn.Module = L.slot(ActivationGroup)   # -> '${activation}'

    schema = MLPSchema

    # The defaults list selects the default entry for the group. L.bind composes
    # ActivationGroup.relu under the slot's package ('activation').
    defaults = L.Defaults(
        L.self_,
        L.bind(MLPSchema.activation, ActivationGroup.relu),
    )

    model = L.call(nn.Sequential, root=True)(
        L.OrderedDict(
            ("input", L.call(nn.Sequential)(
                L.call(nn.Linear)(
                    in_features=L.ref("${schema.dim_in}"),
                    out_features=L.ref("${schema.dim_hidden}"),
                ),
                L.chosen(ActivationGroup, package="activation"),
            )),
            ("hidden", L.call(nn.Sequential, expand_args=True)(
                L.repeat(
                    L.ref("${schema.num_layers}"),
                    L.call(nn.Sequential)(
                        L.call(nn.Linear)(
                            in_features=L.ref("${schema.dim_hidden}"),
                            out_features=L.ref("${schema.dim_hidden}"),
                        ),
                        L.chosen(ActivationGroup, package="activation"),
                    ),
                ),
            )),
            ("output", L.call(nn.Sequential)(
                L.call(nn.Linear)(
                    in_features=L.ref("${schema.dim_hidden}"),
                    out_features=L.ref("${schema.dim_out}"),
                )
            )),
        )
    )

    __all__ = ["model", "schema", "defaults"]
""")

_mlp_path = pathlib.Path(tempfile.mkdtemp()) / "mlp.py"
_mlp_path.write_text(_MLP_TYPED)

# Inspect the schema fragment (the typed dataclass node).
schema_cfg = laco.load(f"{_mlp_path}#schema")
print("--- schema fragment ---")
print(laco.dump(schema_cfg))

Output

--- schema fragment ---
!!python/object:builtins.MLPSchema
activation: {_convert_: all, _target_: torch.nn.ReLU}
dim_hidden: 256
dim_in: 128
dim_out: 128
num_layers: 3

# Load the full config (so the model body can resolve cross-references into
# `schema`) and show the composed activation node — the defaults list selected
# `relu`, which L.chosen pulled into every layer.
cfg = laco.load(str(_mlp_path))
print("--- composed activation (default: relu) ---")
print(laco.dump(cfg.activation))
print("--- instantiated input block (note the resolved ${activation} -> ReLU) ---")
print(laco.instantiate(cfg.model).input)

Output

--- composed activation (default: relu) ---
{_convert_: all, _laco_: 1, _target_: torch.nn.ReLU}

--- instantiated input block (note the resolved ${activation} -> ReLU) ---
Sequential(
  (0): Linear(in_features=128, out_features=256, bias=True)
  (1): ReLU()
)

Section 7: Overriding a Group Selection

The real payoff: swapping the activation is a single CLI-style override. The model config body is never touched: L.chosen(ActivationGroup) resolves to whatever the defaults list selected.

# Default selection: relu
cfg_relu = laco.load(str(_mlp_path))
print("=== relu (default) ===")
print(laco.dump(cfg_relu.activation))

Output

=== relu (default) ===
{_convert_: all, _laco_: 1, _target_: torch.nn.ReLU}

# Switch the activation to gelu. The model body is untouched — overriding the
# composed group node's target propagates through every L.chosen(ActivationGroup)
# reference at once. (We also bump dim_in to show scalar overrides compose too.)
cfg_gelu = laco.load(
    str(_mlp_path),
    "schema.dim_in=64",
    "activation._target_=torch.nn.GELU",   # swap the selected activation entry
)
print("=== gelu (override) ===")
print(laco.dump(cfg_gelu.activation))

Output

=== gelu (override) ===
{_convert_: all, _laco_: 1, _target_: torch.nn.GELU}

Notice that the _target_ in the composed nodes changes from torch.nn.modules.activation.ReLU to torch.nn.modules.activation.GELU everywhere L.chosen(ActivationGroup) appears, including both the input layer and every hidden layer. A single override propagates consistently through the entire model graph.

# Instantiate both and verify the activation types
model_relu = laco.instantiate(laco.load(str(_mlp_path)).model)
model_gelu = laco.instantiate(
    laco.load(str(_mlp_path), "activation._target_=torch.nn.GELU").model
)

def first_activation(sequential):
    """Return the activation module in the input block."""
    return type(list(sequential.input.children())[1]).__name__

print("relu model — input activation:", first_activation(model_relu))
print("gelu model — input activation:", first_activation(model_gelu))

Output

relu model — input activation: ReLU
gelu model — input activation: GELU

Section 8: Group / Slot / Bind Relationships

The table below shows how the four primitives relate at config-authoring time and at Hydra-compose time.

Step	Primitive	Produces
1	`ActivationGroup(L.Group[nn.Module])` with `relu`, `gelu`, `silu` entries	Each entry registered in Hydra's `ConfigStore` under group `activationgroup`
2	`L.bind(MLPSchema.activation, ActivationGroup.relu)`	Validates the entry's group matches the slot, then a `DefaultsBinding` → `{"activationgroup": "relu"}`
3	`MLPSchema.activation: nn.Module = L.slot(ActivationGroup)`	Field default rewritten to `"${activation}"`
4	`L.chosen(ActivationGroup)` in the model body	Same `"${activation}"` interpolation, resolved at instantiation time

At compose time, the defaults-list binding (step 2) fills the schema's slot (step 3), and every L.chosen(ActivationGroup) reference in the model body (step 4) resolves to that same selected entry.

Section 9: `@L.params` vs `L.Group` Side-by-Side

Both approaches can describe configurable components. The table below summarizes the trade-offs so you can choose the right tool for each use case.

# ============================================================
#  @L.params approach                  L.Group approach
# ============================================================

# --- @L.params ---                    # --- L.Group ---
@L.params                               # class ActivationGroup(L.Group[nn.Module]):
class hps_params:                       #     relu = L.call(nn.ReLU)()
    activation = nn.ReLU                #     gelu = L.call(nn.GELU)()
                                        #     silu = L.call(nn.SiLU)()
                                        #
# Override via CLI:                     # Override via CLI:
#   activation=nn.GELU                  #   +activationgroup=gelu
#   (a free-form string — no           #   (validated at config-authoring
#    validation until runtime)          #    time; typos are static errors)

print("@L.params: activation default type:",
      type(hps_params.activation).__name__)
print("L.Group: ActivationGroup entries:",
      [e.name for e in ActivationGroup.entries()])

Output

@L.params: activation default type: str
L.Group: ActivationGroup entries: ['relu', 'gelu', 'silu']

Property	`@L.params`	`L.Group[T]`
Validated choice set	No — any string accepted	Yes — class attrs only
Static type errors	No — runtime crash on typo	Yes — pyright catches typos
IDE autocomplete	Partial — scalars only	Yes — entry attributes
Nested object nodes	No — scalar values only	Yes — full `DictConfig` nodes
Registered in ConfigStore	No	Yes — `ConfigStore.store()`
Hydra group composition	No	Yes — via defaults list
Best for	Scalar hyperparameters (`lr`, `bs`, …)	Swappable components (model, optim, …)

Section 10: Bonus, `L.required[T]()` in Schemas

Sometimes a schema field has no sensible default: it must be explicitly provided by every caller. L.required[T]() marks such fields at the schema level, producing OmegaConf.MISSING as the default value. Attempting to load a config without supplying the required field raises a MissingMandatoryValue error at compose time.

The text-classifier example uses this for vocab_size:

@L.config
class TextClassifierSchema:
    """vocab_size is required — no default."""

    vocab_size:   int       = L.required[int]()
    embed_dim:    int       = 64
    num_classes:  int       = 2
    padding_idx:  int | None = 0

# The field is MISSING at the dataclass level
import dataclasses
from omegaconf import MISSING
vocab_field = dataclasses.fields(TextClassifierSchema)[0]
print(f"vocab_size default: {vocab_field.default!r}")
print(f"Is MISSING: {vocab_field.default is MISSING}")

Output

vocab_size default: '???'
Is MISSING: True

# Loading without providing vocab_size raises a clear error at compose time
from omegaconf import MissingMandatoryValue

try:
    # This will raise because vocab_size has no default
    cfg_missing = laco.load("configs://examples/typed/text_classifier.py#schema")
    # Accessing the missing field triggers the error
    _ = cfg_missing.vocab_size
except MissingMandatoryValue as e:
    print(f"MissingMandatoryValue: {e}")

Output

MissingMandatoryValue: Missing mandatory value: schema.vocab_size
    full_key: schema.vocab_size
    object_type=TextClassifierSchema

# Providing the value via an override resolves it cleanly
cfg_ok = laco.load(
    "configs://examples/typed/text_classifier.py?schema.vocab_size=1000#schema"
)
print(f"vocab_size: {cfg_ok.vocab_size}")
print(f"embed_dim:  {cfg_ok.embed_dim}")

Output

vocab_size: 1000
embed_dim:  64

Summary

Primitive	What it does
`class G(L.Group[T])`	Declares a group of interchangeable config nodes; registers each attribute in Hydra's ConfigStore
`@L.config`	Turns a schema class into a real `dataclass`; resolves `L.slot` fields into `${package}` interpolations
`L.slot(G)`	Field specifier: links a schema field to group `G`; default becomes `"${field_name}"`
`L.bind(slot, entry)`	Produces a `DefaultsBinding` (`{group: name}`); validates group membership at authoring time
`L.Defaults(L.self_, L.bind(...))`	Builds a Hydra defaults list that tells the composer which entry to use for each slot
`L.chosen(G)`	In-body interpolation that resolves to whichever entry was selected for group `G`
`L.required[T]()`	Marks a schema field as having no default (`MISSING`); forces callers to supply a value

When to use @L.params vs L.Group[T]:

Use @L.params for scalar hyperparameters: learning rate, batch size, and epoch count are values you tweak per run.
Use L.Group[T] for swappable components: optimizer family, activation function, loss function, and encoder architecture are objects you swap wholesale.

Next: 08.pipeline-configs.ipynb shows how to wire multiple typed components into a full training bundle with L.Dict, relative imports, and the @L.task entry point.

Nested Configs and Containers

Pipeline Configs