Upload folder using huggingface_hub

ecc81b3 verified 3 months ago

7.3 kB

"""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))

	"""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))