"""The Vision Transformer, and why it is an N-D model at all. ViT is usually described as "images as sequences of patches", which hides the lattice: the patches *are* a grid, and every design choice that distinguishes ViT variants is a choice about how that grid is handled. td.ViT(d_model=192, n_layers=6, image=(32, 32), patch=4) # joint, as published td.ViT(..., method=td.axial_scan) # axial ViT The first is the published model: patchify, flatten the grid to one sequence, attend over all of it. The second is the axial variant — one argument apart, which is the entire point of the library. On an 8×8 patch grid the joint form attends over 64 tokens; the axial form does two passes of 8. Same patches, same parameter count, different method of multidimensionality, and BENCHMARKS.md says which is cheaper where. The kernel family (``td.cafa``, ``td.axial_attention``) is deliberately *not* available here, and the refusal is informative rather than a gap: those methods own every spatial axis themselves and leave the mixer to run along time, so on a time-less patch grid the transformer blocks would be dead weight. A factorized-attention model over a patch grid is ``td.AxialKernel(mixer=None, ...)`` — the kernels are the model. Give the lattice a time axis (video) and the hybrid form applies again. **What this ships and what it does not.** Patch embedding, positional embedding, and the transformer stack over the patch lattice, returning per-patch features ``(B, *grid, d_model)``. No class token, no pooling, no classification head — the same boundary every other model in the library keeps. A head is three lines of caller code and it is the caller's three lines. Concretely:: vit = td.ViT(192, 6, image=(32, 32), patch=4, in_channels=3) head = nn.Linear(192, 10) logits = head(vit(images).mean(dim=(1, 2))) # mean-pool the grid **Positional embedding is where the lattice shows.** ViT learns one embedding per patch position — a table the size of the grid, which cannot transfer to a different image size. The factorized alternative learns one table per *axis* and adds them, which is ``r·A`` parameters instead of ``A^r`` and extends to a new grid size by interpolating one axis at a time. Both are here; factorized is the default because on a lattice it is the natural one, and the published choice is one argument away. """ from __future__ import annotations from collections.abc import Sequence import torch import torch.nn as nn from torch_dimensions.compose import flatten from torch_dimensions.lattice import Lattice from torch_dimensions.mixers.attention import AttentionMixer from torch_dimensions.models.base import LatticeModel __all__ = ["PatchEmbed", "ViT"] class PatchEmbed(nn.Module): """Cut an image into patches and embed each one — image to lattice. ``(B, *image, C)`` in, ``(B, *grid, d_model)`` out, where ``grid[i] = image[i] // patch[i]``. Rank-generic: a 2-D image, a 3-D volume, and a 4-D spatio-temporal block all work, because the patching is a reshape and a linear map rather than a ``Conv2d`` with a rank baked in. Args: image: size of each input axis. patch: patch size per axis; an int applies to every axis. Must divide the image exactly — a partial patch at the edge is a silent crop, and cropping the user's data without saying so is not this module's decision to make. in_channels: channels of the input (3 for RGB, 1 for greyscale). d_model: embedding width per patch. """ def __init__( self, image: Sequence[int], patch: Sequence[int] | int, in_channels: int, d_model: int, ) -> None: super().__init__() self.image = tuple(int(s) for s in image) rank = len(self.image) self.patch = ( (int(patch),) * rank if isinstance(patch, int) else tuple(int(p) for p in patch) ) if len(self.patch) != rank: raise ValueError(f"patch {self.patch} has {len(self.patch)} axes, image has {rank}") bad = [(s, p) for s, p in zip(self.image, self.patch, strict=True) if p < 1 or s % p] if bad: raise ValueError( f"patch size must divide the image exactly; got image={self.image}, " f"patch={self.patch}. A partial patch at the edge would silently crop the " "input — pad or resize before this point, deliberately." ) self.grid = tuple(s // p for s, p in zip(self.image, self.patch, strict=True)) self.in_channels = in_channels self.d_model = d_model self.n_patch_features = in_channels for p in self.patch: self.n_patch_features *= p self.proj = nn.Linear(self.n_patch_features, d_model) def lattice(self, *, names: Sequence[str] | None = None, time: bool = False) -> Lattice: """The patch grid as a lattice — what the model actually operates on.""" return Lattice(shape=self.grid, names=tuple(names) if names else None, time=time) def forward(self, x: torch.Tensor) -> torch.Tensor: rank = len(self.image) if x.ndim != rank + 2: raise ValueError( f"expected a {rank + 2}-D tensor (B, *{self.image}, {self.in_channels}); " f"got shape {tuple(x.shape)}" ) if tuple(x.shape[1:-1]) != self.image: raise ValueError(f"expected image dims {self.image}, got {tuple(x.shape[1:-1])}") if x.shape[-1] != self.in_channels: raise ValueError(f"expected {self.in_channels} channels, got {x.shape[-1]}") b = x.shape[0] # Split each image axis into (grid, patch), then move every patch axis # next to the channel axis and flatten them together. Written as one # rank-generic reshape+permute rather than einops or a per-rank table, # for the same reason the fold is: a rank-4 case must not need new code. split: list[int] = [b] for g, p in zip(self.grid, self.patch, strict=True): split += [g, p] split.append(self.in_channels) h = x.reshape(*split) grid_dims = [1 + 2 * i for i in range(rank)] patch_dims = [2 + 2 * i for i in range(rank)] h = h.permute(0, *grid_dims, *patch_dims, h.ndim - 1).contiguous() h = h.reshape(b, *self.grid, self.n_patch_features) return self.proj(h) def extra_repr(self) -> str: return ( f"image={self.image}, patch={self.patch}, grid={self.grid}, " f"in_channels={self.in_channels}, d_model={self.d_model}" ) class _PosEmbed(nn.Module): """Learned positional embedding over a patch grid. ``factorized`` learns one table per axis and adds them (``r·A`` parameters); ``full`` learns one per cell (``A^r``), which is what ViT publishes. Factorized is the default: on a lattice it is the natural parameterization, it is what makes a 3-D or 4-D grid affordable, and the axial models in this library already assume per-axis structure everywhere else. """ def __init__(self, grid: tuple[int, ...], d_model: int, kind: str) -> None: super().__init__() if kind not in ("factorized", "full", "none"): raise ValueError(f"pos_embed must be factorized|full|none; got {kind!r}") self.kind = kind self.grid = grid if kind == "factorized": self.tables = nn.ParameterList(nn.Parameter(torch.zeros(n, d_model)) for n in grid) elif kind == "full": self.tables = nn.ParameterList([nn.Parameter(torch.zeros(*grid, d_model))]) else: self.tables = nn.ParameterList() for t in self.tables: nn.init.trunc_normal_(t, std=0.02) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.kind == "none": return x if self.kind == "full": return x + self.tables[0] rank = len(self.grid) for axis, table in enumerate(self.tables): # Broadcast the axis's table along every other grid axis. shape = [1] * rank shape[axis] = self.grid[axis] x = x + table.reshape(*shape, -1) return x def extra_repr(self) -> str: n = sum(p.numel() for p in self.tables) return f"{self.kind}, grid={self.grid}, {n} parameters" class ViT(LatticeModel): """Vision Transformer over a patch lattice, at any rank. Args: d_model: embedding width. n_layers: transformer blocks. image: input size per axis — ``(32, 32)`` for CIFAR, ``(16, 64, 64)`` for a volume. patch: patch size per axis, or one int for all. in_channels: input channels. pos_embed: ``"factorized"`` (default), ``"full"`` (ViT's own), or ``"none"``. method: the method of multidimensionality. Defaults to :func:`~torch_dimensions.flatten` — attention over all patches at once, which is what ViT does. ``td.axial_scan`` gives the axial variant. The kernel family needs a time axis; see the module docstring. mixer_kwargs: forwarded to :class:`~torch_dimensions.mixers.attention.AttentionMixer`, e.g. ``{"n_heads": 12}``. Returns per-patch features ``(B, *grid, d_model)``. Pool and classify yourself; see the module docstring. """ _mixer = AttentionMixer def __init__( self, d_model: int, n_layers: int = 1, *, image: Sequence[int], patch: Sequence[int] | int = 16, in_channels: int = 3, pos_embed: str = "factorized", names: Sequence[str] | None = None, n_heads: int = 4, **kw, ) -> None: # The lattice is derived, so accepting one would be accepting a second # answer to a question already answered. A checkpoint records # `lattice: None` for exactly this reason, and rebuilding hands it back. if kw.pop("lattice", None) is not None: raise ValueError( "ViT builds its lattice from `image` and `patch`; passing `lattice` too " "would let the two disagree. Pass image/patch, or use td.Transformer " "directly if you already have a lattice of tokens." ) embed = PatchEmbed(image, patch, in_channels, d_model) lat = embed.lattice(names=names) # Only default when the caller named neither spelling: setting # `nd_method` unconditionally would collide with a caller's `method=` # and the base class refuses both, as it should. if "method" not in kw and "nd_method" not in kw: kw["nd_method"] = flatten mixer_kwargs = {"n_heads": n_heads, **kw.pop("mixer_kwargs", {})} super().__init__(d_model, n_layers, lat, mixer_kwargs=mixer_kwargs, **kw) # Registered after super().__init__ so the base class's parameter # accounting and the recorded config are unaffected by their presence. self.patch_embed = embed self.pos = _PosEmbed(embed.grid, d_model, pos_embed) self.config.update( { "image": list(embed.image), "patch": list(embed.patch), "in_channels": in_channels, "pos_embed": pos_embed, "n_heads": n_heads, # The lattice is derived from image/patch, so recording it too # would let a checkpoint hold two answers to one question. "lattice": None, } ) @property def grid(self) -> tuple[int, ...]: """Patch-grid shape — the lattice the transformer runs over.""" return self.patch_embed.grid def forward(self, x: torch.Tensor) -> torch.Tensor: """``(B, *image, C)`` in, ``(B, *grid, d_model)`` out.""" return self.nd(self.pos(self.patch_embed(x)))