Hyperparameters and Interpolation

import laco
import laco.language as L
import torch.nn as nn

# Hard-coded dimensions — fragile!
layer1 = L.call(nn.Linear)(in_features=128, out_features=256)
layer2 = L.call(nn.Linear)(in_features=256, out_features=256)
layer3 = L.call(nn.Linear)(in_features=256, out_features=64)

# To change the hidden size from 256 → 512, you must find every "256"
# and decide whether it is a hidden dimension or something else.
# Mistakes are silent — nothing checks that layer2.in == layer1.out.
print(laco.dump(layer2))

@L.params
class hps:
    dim_in:     int = 128
    dim_out:    int = 64
    dim_hidden: int = 256
    num_layers: int = 3

# Attribute access returns interpolation strings at runtime:
print(hps.dim_in)      # "${hps.dim_in}"
print(hps.dim_out)     # "${hps.dim_out}"
print(hps.dim_hidden)  # "${hps.dim_hidden}"
print(type(hps.dim_in))  # <class 'str'>

layer1_cfg = L.call(nn.Linear)(in_features=hps.dim_in,     out_features=hps.dim_hidden)
layer2_cfg = L.call(nn.Linear)(in_features=hps.dim_hidden, out_features=hps.dim_hidden)
layer3_cfg = L.call(nn.Linear)(in_features=hps.dim_hidden, out_features=hps.dim_out)

# Each layer node holds the interpolation strings, e.g. in_features: ${hps.dim_in}.
# An interpolation can only resolve when the *full* config tree also contains the
# hps node with the actual integers. On its own a fragment has no hps sibling, so
# dumping it would raise InterpolationKeyError. We assemble that tree by hand here;
# laco.load does exactly this for a config file (see Section 4). Calling hps() on
# the @L.params block returns its plain {field: value} dict.
tree = {"hps": hps(), "layer1": layer1_cfg, "layer2": layer2_cfg, "layer3": layer3_cfg}

print(laco.dump(tree))

@L.params
class dims:
    width: int   = 512
    depth: int   = 6
    drop:  float = 0.1

# What Python holds at runtime:
print("dims.width  (runtime) =", repr(dims.width))   # '${dims.width}'
print("dims.depth  (runtime) =", repr(dims.depth))   # '${dims.depth}'
print("dims.drop   (runtime) =", repr(dims.drop))    # '${dims.drop}'

# The IDE / pyright reports:
#   dims.width : int
#   dims.depth : int
#   dims.drop  : float
# ... which is a lie, but a useful one.
print()
print("Runtime type of dims.width:", type(dims.width))  # str

@L.params class hps
    → at import time: ParamsWrapper(hps)

hps.dim_hidden  (attribute access)
    → "${hps.dim_hidden}"  (str)

L.call(nn.Linear)(in_features=hps.dim_hidden, ...)
    → DictConfig { _target_: "...", in_features: "${hps.dim_hidden}", ... }

laco.instantiate(full_cfg)  # full_cfg contains the hps node with actual values
    → OmegaConf resolves "${hps.dim_hidden}" → 256
    → nn.Linear(in_features=256, ...)

@L.params
class hps:
    in_features:   int   = 8
    out_features:  int   = 1
    learning_rate: float = 1e-2

# root=True marks this node as the root of a config tree (it is a typing/intent
# marker; it does not bundle the hps node by itself). For the ${hps.*} references
# to resolve, the hps node must live alongside the model in one tree -- which is
# precisely what laco.load assembles from a file's exported names. We mirror that
# here so the config is self-contained and dumpable.
model_cfg = L.call(nn.Linear, root=True)(
    in_features  = hps.in_features,
    out_features = hps.out_features,
)

full_cfg = {"hps": hps(), "model": model_cfg}

print("--- model config YAML ---")
print(laco.dump(full_cfg))

# Load from the real file with an override:
cfg = laco.load(
    "configs://examples/linear_regression.py",
    "hps.in_features=32",
    key="model"
)
built = laco.instantiate(cfg)
print(built)  # Linear(in_features=32, out_features=1, bias=True)

# Explicitly typed interpolation reference
hidden_ref: int = L.ref("${hps.dim_hidden}")
print(hidden_ref)        # "${hps.dim_hidden}"
print(type(hidden_ref))  # <class 'str'>  (typed as int by the IDE)

# Cross-file reference (e.g., referencing an encoder's output dim
# from a decoder config that lives in a different file)
encoder_out_ref: int = L.ref("${encoder.out_dim}")
print(encoder_out_ref)

@L.params
class model_hps:
    hidden: int = 512
    num_heads: int = 8

# These all produce interpolation strings, not computed values:
print("sum    :", L.r.sum(model_hps.hidden, 64))
print("div    :", L.r.div(model_hps.hidden, 2))
print("pow    :", L.r.pow(2, model_hps.num_heads))
print("head_dim:", L.r.div(model_hps.hidden, model_hps.num_heads))

# Practical use: derive head_dim from hidden and num_heads
# so changing either one automatically updates the other.
head_dim_cfg = L.call(nn.Linear)(
    in_features  = L.r.div(model_hps.hidden, model_hps.num_heads),
    out_features = model_hps.hidden,
)

# As before, the resolver string ${div:${model_hps.hidden},${model_hps.num_heads}}
# only resolves when the model_hps node is present in the same tree, so we bundle
# it in (just like laco.load does for a config file).
full_cfg = {"model_hps": model_hps(), "proj": head_dim_cfg}

print("--- head projection config ---")
print(laco.dump(full_cfg))

# ---- Syntax 1: URL query string ----
# Append ?key=value&key2=value2 and optionally #fragment to select a sub-key.
cfg_url = laco.load(
    "configs://examples/linear_regression.py?hps.in_features=32&hps.out_features=4#model"
)
print("URL override — model:")
print(laco.dump(cfg_url))

# ---- Syntax 2: Positional override strings ----
# Pass extra strings after the path; each is "key=value".
cfg_pos = laco.load(
    "configs://examples/linear_regression.py",
    "hps.in_features=32",
    "hps.out_features=4",
    key="model",
)
print("Positional override — model:")
print(laco.dump(cfg_pos))

# Side-by-side: default vs. overridden
cfg_default  = laco.load("configs://examples/linear_regression.py")
cfg_override = laco.load(
    "configs://examples/linear_regression.py",
    "hps.in_features=64",
    "hps.learning_rate=5e-4",
)

print("=== default hps ===")
print(laco.dump(cfg_default["hps"]))

print("=== overridden hps ===")
print(laco.dump(cfg_override["hps"]))

# === Part 1: Hyperparameters ===
# Note: activation is typed as type[nn.Module] — laco can store a class object
# as a config value. The YAML representation is a !!python/name: tag.
@L.params
class mlp_hps:
    dim_in:     int           = 128
    dim_out:    int           = 128
    dim_hidden: int           = 256
    num_layers: int           = 3
    activation: type[nn.Module] = nn.ReLU

print("hps.dim_in  :", mlp_hps.dim_in)     # "${mlp_hps.dim_in}"
print("hps.activation:", mlp_hps.activation)  # "${mlp_hps.activation}"

# === Part 2: make_mlp factory function ===
# The function takes the resolved values as arguments — not hps attributes.
# This makes it usable both from the config system (where hps.* are
# interpolation strings) and from plain Python code (where you pass integers).

def make_mlp(*, dim_in, dim_out, dim_hidden, num_layers, activation):
    return L.call(nn.Sequential, root=True)(
        # L.OrderedDict groups named layers into a DictConfig.
        # Each entry is a ("name", config_node) tuple.
        L.OrderedDict(
            (
                "input",
                # The input block: Linear → activation
                L.call(nn.Sequential)(
                    L.call(nn.Linear)(in_features=dim_in, out_features=dim_hidden),
                    L.call(activation)(),  # activation is a class reference
                ),
            ),
            (
                "hidden",
                # L.repeat(n, src) produces n deep-copies of src in a list.
                # expand_args=True unpacks the list as positional args to Sequential.
                L.call(nn.Sequential, expand_args=True)(
                    L.repeat(
                        num_layers,
                        L.call(nn.Sequential)(
                            L.call(nn.Linear)(
                                in_features=dim_hidden,
                                out_features=dim_hidden,
                            ),
                            L.call(activation)(),
                        ),
                    ),
                ),
            ),
            (
                "output",
                L.call(nn.Sequential)(
                    L.call(nn.Linear)(in_features=dim_hidden, out_features=dim_out),
                ),
            ),
        )
    )

# === Part 3: module-level model symbol ===
# Called with hps.* — all values are interpolation strings here.
# laco.load exports this symbol and assembles the full tree with the hps node,
# so OmegaConf can resolve the interpolations at instantiate time.
model = make_mlp(
    dim_in     = mlp_hps.dim_in,
    dim_out    = mlp_hps.dim_out,
    dim_hidden = mlp_hps.dim_hidden,
    num_layers = mlp_hps.num_layers,
    activation = mlp_hps.activation,
)

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

# Load the real mlp.py config and instantiate it:
mlp_cfg = laco.load("configs://examples/mlp.py#model")
mlp_model = laco.instantiate(mlp_cfg)
print(mlp_model)
print()
# Count parameters:
total_params = sum(p.numel() for p in mlp_model.parameters())
print(f"Total parameters: {total_params:,}")

# Sweep dim_hidden and num_layers without touching the file:
for dim_hidden, num_layers in [(128, 2), (256, 3), (512, 4)]:
    cfg = laco.load(
        "configs://examples/mlp.py",
        f"hps.dim_hidden={dim_hidden}",
        f"hps.num_layers={num_layers}",
        key="model"
    )
    built = laco.instantiate(cfg)
    params = sum(p.numel() for p in built.parameters())
    print(f"  dim_hidden={dim_hidden:3d}, num_layers={num_layers} → {params:>7,} params")