"""The shared shape of every model in the library. A model here is nothing but: an optional input projection, a 1-D mixer type, and an ``nd_method`` deciding how that mixer covers the lattice. One class holds that recipe; ``LSTM``, ``GRU``, ``S4D``, and ``Mamba`` differ only in ``_mixer``. This is the design's central claim made literal — an N-D model is a 1-D mixer plus a plan for sweeping it. """ from __future__ import annotations import warnings from collections.abc import Callable from functools import partial from typing import cast import torch import torch.nn as nn from torch_dimensions.compose import axial_scan, resolve_nd_method from torch_dimensions.lattice import AxisSpec, Lattice from torch_dimensions.plan import ScanPlan __all__ = ["LatticeModel"] class LatticeModel(nn.Module): """Base for the model family. Subclasses set ``_mixer``. With no lattice this is an ordinary sequence model — a lattice with no spatial axes has an identity permutation, so the 1-D case is the N-D case with nothing to fold. """ _mixer: type[nn.Module] def __init__( self, d_model: int, n_layers: int = 1, lattice: Lattice | None = None, *, d_input: int | None = None, nd_method: str | Callable[..., nn.Module] = axial_scan, method: str | Callable[..., nn.Module] | None = None, plan: ScanPlan | None = None, bidirectional: bool | AxisSpec | list[AxisSpec] = False, dropout: float = 0.0, chunk: int | None = None, mixer_kwargs: dict | None = None, mixer: type[nn.Module] | None = None, **method_kwargs, ) -> None: super().__init__() # `method` is the short spelling of `nd_method` — the method of # multidimensionality. Both name the same thing; giving both is a # contradiction waiting to happen and is refused. if method is not None: if nd_method is not axial_scan: raise ValueError("pass either `method` or `nd_method`, not both") nd_method = method # No lattice means a single dynamic axis: an ordinary sequence. self.lattice = lattice if lattice is not None else Lattice(shape=(), time=True) if plan is None: plan = ScanPlan.cyclic(self.lattice.axis_names, n_layers, bidirectional=bidirectional) else: if bidirectional is not False: raise ValueError("pass either `plan` or `bidirectional`, not both") # A plan *is* the layer schedule, so it fixes the depth, and it # wins. But winning *silently* over a disagreeing n_layers would # ship a model shallower (or deeper) than requested — the same # silent-downgrade failure the schedule machinery exists to make # loud. A warning rather than an error because generic builders # legitimately fill n_layers unconditionally and add a plan only # sometimes; n_layers=1 is the default and passes untouched. if n_layers != 1 and n_layers != len(plan): warnings.warn( f"n_layers={n_layers} is ignored: the given plan has {len(plan)} steps " "and a plan determines the depth", UserWarning, stacklevel=2, ) # An input projection only when the data is not already d_model wide. # Without it every caller writes the same nn.Linear, which is friction # for no gain in purity. self.in_proj = nn.Linear(d_input, d_model) if d_input is not None else nn.Identity() # `mixer=` substitutes the 1-D operator without touching anything else, # which is what makes a model debuggable: swapping in # `td.testing.Recorder` answers "which axis did layer 3 sweep" on the # real model rather than on a reconstruction of it. The class's own # `_mixer` stays the default and the recorded config still names the # class, so a substituted model is visibly not the stock one. mixer_cls = mixer if mixer is not None else self._mixer # A substituted mixer cannot be written into the recipe — a class is # not serializable — so a checkpoint would silently rebuild with the # stock one and hand back a *different model* that loads without # complaint. `save()` refuses instead; see config._checkpoint_header. self._substituted_mixer = None if mixer is None else mixer.__name__ self.nd = resolve_nd_method(nd_method)( mixer=partial(mixer_cls, d_model, **(mixer_kwargs or {})), plan=plan, lattice=self.lattice, d_model=d_model, dropout=dropout, chunk=chunk, **method_kwargs, ) # The construction recipe, recorded so a checkpoint can rebuild this # model without the user re-specifying anything (td.save / td.load). # Plain JSON-able types throughout — the validity mask becomes a # nested list — so a config can also live in YAML unchanged. n_layers # is recorded as the plan's true depth: the plan is the schedule, and # a recipe that could disagree with itself would not be a recipe. from torch_dimensions.config import lattice_to_dict, nd_method_name self.config: dict = { "d_model": d_model, "n_layers": len(plan), "lattice": None if lattice is None else lattice_to_dict(self.lattice), "plan": plan.to_dict(), "nd_method": nd_method_name(nd_method), "d_input": d_input, "dropout": dropout, "chunk": chunk, "mixer_kwargs": dict(mixer_kwargs) if mixer_kwargs else {}, **method_kwargs, } @property def plan(self) -> ScanPlan: return cast(ScanPlan, self.nd.plan) def save(self, path) -> None: """Write this model — architecture and weights — to one checkpoint file that :func:`torch_dimensions.load` can rebuild it from.""" from torch_dimensions.config import save save(self, path) def receptive_field(self) -> dict[str, dict[str, object]]: """How far along each axis this model can actually see. ``inf`` for mixers that span their axis in one layer (RNNs, SSMs, attention); a finite number for convolutions, where it is a real constraint worth checking before training rather than after. See :func:`torch_dimensions.receptive_field`. """ from torch_dimensions.mixers.conv import axis_receptive_field return axis_receptive_field(self) def to_spec(self) -> dict: """A JSON-able description of this model's N-D architecture. Derived without a forward pass, so it can be taken before any data exists. See VIEWER.md. """ from torch_dimensions.spec import scan_model_spec return scan_model_spec(self) def forward(self, x: torch.Tensor) -> torch.Tensor: """``(B, [T,] *shape, d_input or d_model)`` in, ``d_model`` out. With no lattice that is just ``(B, T, ...)``. """ return self.nd(self.in_proj(x))