stackformers 3.8.0__py3-none-any.whl

stackformers/positional/config.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class YaRNConfig(BaseModel):
+    """NTK-by-parts RoPE frequency scaling (Peng et al., 2023).
+
+    Extends RoPE to longer contexts by scaling low-frequency dimensions
+    linearly while leaving high-frequency dimensions unchanged.  The
+    inv_freq buffer is modified once at construction time so forward()
+    stays export-traceable.
+
+    scale:               target_max_seq_len / original_max_seq_len
+    original_max_seq_len: context length the base model was trained on
+    beta_fast:            high-freq threshold (wavelength = original / beta_fast)
+    beta_slow:            low-freq threshold  (wavelength = original / beta_slow)
+    """
+
+    scale: float = Field(gt=1.0)
+    original_max_seq_len: int = Field(gt=0)
+    beta_fast: float = Field(default=32.0, gt=0.0)
+    beta_slow: float = Field(default=1.0, gt=0.0)
+
+    @model_validator(mode="after")
+    def _check_beta_ordering(self) -> YaRNConfig:
+        if self.beta_slow >= self.beta_fast:
+            raise ValueError(
+                f"beta_slow ({self.beta_slow}) must be less than beta_fast ({self.beta_fast}). "
+                "Swapped values invert the high/low-frequency partition and silently break scaling."
+            )
+        return self
+
+
+class RoPE1DConfig(BaseModel):
+    kind: Literal["rope1d"] = "rope1d"
+    dim_head: int = Field(gt=0)
+    base: int = Field(default=10_000, gt=0)
+    yarn: YaRNConfig | None = None
+
+
+class RoPE2DConfig(BaseModel):
+    kind: Literal["rope2d"] = "rope2d"
+    dim_head: int = Field(gt=0)
+    base: int = Field(default=10_000, gt=0)
+
+
+class NoPosEncodingConfig(BaseModel):
+    kind: Literal["none"] = "none"
+
+
+class LearnedPosEncodingConfig(BaseModel):
+    """Config for learned absolute position embeddings (Vaswani et al., 2017)."""
+
+    kind: Literal["learned"] = "learned"
+    dim_head: int = Field(gt=0)
+    max_seq_len: int = Field(gt=0)
+
+
+PosEncodingConfig = Annotated[
+    RoPE1DConfig | RoPE2DConfig | NoPosEncodingConfig | LearnedPosEncodingConfig,
+    Field(discriminator="kind"),
+]

stackformers/positional/factory.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from stackformers.positional.config import (
+    LearnedPosEncodingConfig,
+    NoPosEncodingConfig,
+    PosEncodingConfig,
+    RoPE1DConfig,
+    RoPE2DConfig,
+)
+from stackformers.positional.learned import LearnedPosEncoding
+from stackformers.positional.none import NoPosEncoding
+from stackformers.positional.protocols import PosEncoding
+from stackformers.positional.rope1d import RotaryEmbedding1D
+from stackformers.positional.rope2d import RotaryEmbedding2D
+
+
+def build_pos_encoding(config: PosEncodingConfig) -> PosEncoding:
+    match config:
+        case RoPE1DConfig():
+            return RotaryEmbedding1D(config)
+        case RoPE2DConfig():
+            return RotaryEmbedding2D(config)
+        case NoPosEncodingConfig():
+            return NoPosEncoding(config)
+        case LearnedPosEncodingConfig():
+            return LearnedPosEncoding(config)
+        case _:
+            raise AssertionError(f"Unhandled pos encoding config: {type(config)}")

stackformers/positional/learned.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.positional.config import LearnedPosEncodingConfig
+
+
+class LearnedPosEncoding(nn.Module):
+    """Learned absolute position embeddings added to query and key tensors.
+
+    Each position index is mapped to a learned dh-dimensional vector, which is
+    added to q and k after the head projection.  Requires integer position indices
+    in positions[..., 0]; out-of-range indices raise an embedding error.
+
+    Section 3.5 of Vaswani et al., 2017 describes the learned variant alongside
+    sinusoidal encodings — https://arxiv.org/abs/1706.03762
+    """
+
+    def __init__(self, config: LearnedPosEncodingConfig) -> None:
+        super().__init__()
+        self.emb = nn.Embedding(config.max_seq_len, config.dim_head)
+        nn.init.normal_(self.emb.weight, std=0.02)
+
+    def _encode(
+        self, q: Tensor, k: Tensor, q_positions: Tensor, k_positions: Tensor
+    ) -> tuple[Tensor, Tensor]:
+        q_idx = q_positions[..., 0].long()
+        k_idx = k_positions[..., 0].long()
+        return q + self.emb(q_idx).unsqueeze(1), k + self.emb(k_idx).unsqueeze(1)
+
+    def forward_padded(
+        self,
+        q: Float[Tensor, "b h n dh"],
+        k: Float[Tensor, "b h s dh"],
+        q_positions: Float[Tensor, "b n c"],
+        k_positions: Float[Tensor, "b s c"],
+    ) -> tuple[Float[Tensor, "b h n dh"], Float[Tensor, "b h s dh"]]:
+        return self._encode(q, k, q_positions, k_positions)
+
+    def forward_packed(
+        self,
+        q: Float[Tensor, "nt h dh"],
+        k: Float[Tensor, "nt h dh"],
+        q_positions: Float[Tensor, "nt c"],
+        k_positions: Float[Tensor, "nt c"],
+    ) -> tuple[Float[Tensor, "nt h dh"], Float[Tensor, "nt h dh"]]:
+        return self._encode(q, k, q_positions, k_positions)

stackformers/positional/none.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.positional.config import NoPosEncodingConfig
+
+
+class NoPosEncoding(nn.Module):
+    """Null object for PosEncoding — passes q, k unchanged regardless of layout."""
+
+    def __init__(self, _config: NoPosEncodingConfig = NoPosEncodingConfig()) -> None:
+        super().__init__()
+
+    def forward_padded(
+        self,
+        q: Float[Tensor, "b h n dh"],
+        k: Float[Tensor, "b h s dh"],
+        q_positions: Float[Tensor, "b n c"],
+        k_positions: Float[Tensor, "b s c"],
+    ) -> tuple[Float[Tensor, "b h n dh"], Float[Tensor, "b h s dh"]]:
+        return q, k
+
+    def forward_packed(
+        self,
+        q: Float[Tensor, "nt h dh"],
+        k: Float[Tensor, "nt h dh"],
+        q_positions: Float[Tensor, "nt c"],
+        k_positions: Float[Tensor, "nt c"],
+    ) -> tuple[Float[Tensor, "nt h dh"], Float[Tensor, "nt h dh"]]:
+        return q, k

stackformers/positional/protocols.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Float
+from torch import Tensor
+
+
+@runtime_checkable
+class PosEncoding(Protocol):
+    """Apply positional encoding to projected query and key tensors.
+
+    Separate methods per layout so implementations contain no dispatch logic
+    and callers (which already know their layout) pick the right path directly.
+    Null implementation: NoPosEncoding.
+    """
+
+    def forward_padded(
+        self,
+        q: Float[Tensor, "b h n dh"],
+        k: Float[Tensor, "b h s dh"],
+        q_positions: Float[Tensor, "b n c"],
+        k_positions: Float[Tensor, "b s c"],
+    ) -> tuple[Float[Tensor, "b h n dh"], Float[Tensor, "b h s dh"]]: ...
+
+    def forward_packed(
+        self,
+        q: Float[Tensor, "nt h dh"],
+        k: Float[Tensor, "nt h dh"],
+        q_positions: Float[Tensor, "nt c"],
+        k_positions: Float[Tensor, "nt c"],
+    ) -> tuple[Float[Tensor, "nt h dh"], Float[Tensor, "nt h dh"]]: ...

stackformers/positional/rope1d.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.positional.config import RoPE1DConfig, YaRNConfig
+
+
+def _rotate_half(x: Tensor) -> Tensor:
+    """Split x into [x1 | x2] halves and return [-x2 | x1]."""
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rope_padded_unbatched(t: Tensor, freqs: Tensor) -> Tensor:
+    """t: b h n dh, freqs: n dh — (n dh) broadcasts over (b h n dh) without unsqueeze."""
+    cos = freqs.cos()
+    sin = freqs.sin()
+    return t * cos + _rotate_half(t) * sin
+
+
+def _apply_rope(t: Tensor, freqs: Tensor) -> Tensor:
+    """Works for padded (b h n dh)+(b n dh) and packed (nt h dh)+(nt dh).
+
+    unsqueeze(1) inserts the head dim in both cases.
+    Cast to t.dtype so float32 freqs don't upcast a float16 input.
+    """
+    cos = freqs.cos().to(dtype=t.dtype).unsqueeze(1)
+    sin = freqs.sin().to(dtype=t.dtype).unsqueeze(1)
+    return t * cos + _rotate_half(t) * sin
+
+
+def _yarn_inv_freq(
+    inv_freq: Tensor,
+    cfg: YaRNConfig,
+) -> Tensor:
+    """Return YaRN-scaled inv_freq (Peng et al., 2023, NTK-by-parts)."""
+    wavelen = 2.0 * math.pi / inv_freq
+    low_wavelen = cfg.original_max_seq_len / cfg.beta_slow
+    high_wavelen = cfg.original_max_seq_len / cfg.beta_fast
+    smooth = (
+        (cfg.original_max_seq_len / wavelen - cfg.beta_slow) / (cfg.beta_fast - cfg.beta_slow)
+    ).clamp(0.0, 1.0)
+    scaled = inv_freq / cfg.scale
+    blended = smooth * inv_freq + (1.0 - smooth) * scaled
+    return torch.where(
+        wavelen < high_wavelen, inv_freq, torch.where(wavelen > low_wavelen, scaled, blended)
+    )
+
+
+class RotaryEmbedding1D(nn.Module):
+    """1-D Rotary Position Embedding (Su et al., 2021).
+
+    Optionally accepts YaRNConfig for extended context via NTK-by-parts scaling.
+    """
+
+    def __init__(self, config: RoPE1DConfig) -> None:
+        super().__init__()
+        assert config.dim_head % 2 == 0, "dim_head must be even for RoPE"
+        inv_freq = 1.0 / (
+            config.base ** (torch.arange(0, config.dim_head, 2).float() / config.dim_head)
+        )
+        if config.yarn is not None:
+            inv_freq = _yarn_inv_freq(inv_freq, config.yarn)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def _freqs_from_positions(self, positions: Tensor) -> Tensor:
+        """positions: (..., c) → freqs: (..., dh), uses first coordinate only.
+
+        float32 cast ensures half-precision inputs don't lose precision in the outer product.
+        """
+        inv: Tensor = self.inv_freq  # type: ignore[assignment]
+        pos = positions[..., 0].to(dtype=torch.float32)
+        freqs = pos.unsqueeze(-1) * inv.float()
+        return torch.cat([freqs, freqs], dim=-1)
+
+    def _encode(
+        self,
+        q: Tensor,
+        k: Tensor,
+        q_positions: Tensor,
+        k_positions: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        return (
+            _apply_rope(q, self._freqs_from_positions(q_positions)),
+            _apply_rope(k, self._freqs_from_positions(k_positions)),
+        )
+
+    def forward_padded(
+        self,
+        q: Float[Tensor, "b h n dh"],
+        k: Float[Tensor, "b h s dh"],
+        q_positions: Float[Tensor, "b n c"],
+        k_positions: Float[Tensor, "b s c"],
+    ) -> tuple[Float[Tensor, "b h n dh"], Float[Tensor, "b h s dh"]]:
+        return self._encode(q, k, q_positions, k_positions)
+
+    def forward_packed(
+        self,
+        q: Float[Tensor, "nt h dh"],
+        k: Float[Tensor, "nt h dh"],
+        q_positions: Float[Tensor, "nt c"],
+        k_positions: Float[Tensor, "nt c"],
+    ) -> tuple[Float[Tensor, "nt h dh"], Float[Tensor, "nt h dh"]]:
+        return self._encode(q, k, q_positions, k_positions)

stackformers/positional/rope2d.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.positional.config import RoPE2DConfig
+from stackformers.positional.rope1d import _apply_rope
+
+
+class RotaryEmbedding2D(nn.Module):
+    """2-D Rotary Position Embedding for (row, col) positions.
+
+    Conforms to PosEncoding: positions must have c=2 (row, col per token).
+    Splits dim_head evenly: first half encodes row, second half column.
+    """
+
+    def __init__(self, config: RoPE2DConfig) -> None:
+        super().__init__()
+        assert config.dim_head % 4 == 0, "dim_head must be divisible by 4 for 2-D RoPE"
+        half_dh = config.dim_head // 2
+        inv_freq = 1.0 / (config.base ** (torch.arange(0, half_dh, 2).float() / half_dh))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def _freqs_from_positions(self, positions: Tensor) -> Tensor:
+        """positions: (..., 2) → freqs: (..., dh)"""
+        inv: Tensor = self.inv_freq  # type: ignore[assignment]
+        row = positions[..., 0].to(dtype=inv.dtype).unsqueeze(-1) * inv  # (..., dh//4)
+        col = positions[..., 1].to(dtype=inv.dtype).unsqueeze(-1) * inv
+        half = lambda f: torch.cat([f, f], dim=-1)  # noqa: E731
+        return torch.cat([half(row), half(col)], dim=-1)  # (..., dh)
+
+    def _encode(
+        self, q: Tensor, k: Tensor, q_positions: Tensor, k_positions: Tensor
+    ) -> tuple[Tensor, Tensor]:
+        return (
+            _apply_rope(q, self._freqs_from_positions(q_positions)),
+            _apply_rope(k, self._freqs_from_positions(k_positions)),
+        )
+
+    def forward_padded(
+        self,
+        q: Float[Tensor, "b h n dh"],
+        k: Float[Tensor, "b h s dh"],
+        q_positions: Float[Tensor, "b n c"],
+        k_positions: Float[Tensor, "b s c"],
+    ) -> tuple[Float[Tensor, "b h n dh"], Float[Tensor, "b h s dh"]]:
+        return self._encode(q, k, q_positions, k_positions)
+
+    def forward_packed(
+        self,
+        q: Float[Tensor, "nt h dh"],
+        k: Float[Tensor, "nt h dh"],
+        q_positions: Float[Tensor, "nt c"],
+        k_positions: Float[Tensor, "nt c"],
+    ) -> tuple[Float[Tensor, "nt h dh"], Float[Tensor, "nt h dh"]]:
+        return self._encode(q, k, q_positions, k_positions)

stackformers/presets/README.md

@@ -0,0 +1,63 @@
+# presets
+
+Opinionated `nn.Module` subclasses that wire up the building blocks with fixed structural choices. The caller controls all hyperparameters through a typed config — nothing is hardcoded.
+
+For custom wiring, use the building blocks directly.
+
+## Class hierarchy
+
+Each preset is split into two levels:
+
+- **`*Base[C]`** — abstract generic class parameterised over any config type `C`. Defines `__init__` (wires the stack) and `forward`, declares `build_layers` and `build_norm` as abstract. No assumptions about `C`.
+- **`*`** (concrete) — extends `*Base[ConcreteConfig]`, fills in all `build_*` defaults, and exposes additional overridable hooks (`build_pos_encoding`, `build_attn_bias`, `build_ff`, …).
+
+## Presets
+
+**`TransformerEncoder`** — self-attn → ff per layer. Set `causal=True` in `attn` for GPT-style autoregressive decoding. Set `window_size` for sliding-window (O(n·w)) attention.
+
+**`TransformerDecoder`** — causal self-attn → cross-attn → ff per layer. Cross-attention always uses no positional encoding.
+
+**`CrossAttender`** — cross-attn → ff per layer. `x` is a fixed set of queries with no self-attention. Useful for Perceiver-style or slot-attention architectures.
+
+All presets accept `PaddedInput` or `PackedInput` interchangeably.
+
+## Customising a collaborator (subclass the concrete preset)
+
+Override one or more `build_*` methods; the rest use their defaults:
+
+```python
+class ALiBiEncoder(TransformerEncoder):
+    def build_attn_bias(self, _: TransformerEncoderConfig) -> AttnBias:
+        return ALiBi(self.config.attn.heads)
+```
+
+## Using a fully custom config (subclass the abstract base)
+
+Implement `build_layers` and `build_norm`; the base `__init__` calls them and nothing else:
+
+```python
+class MyConfig(BaseModel):
+    dim: int
+    heads: int
+    num_layers: int
+
+class MyEncoder(TransformerEncoderBase[MyConfig]):
+    def build_layers(self, config: MyConfig) -> list[TransformerLayer]:
+        pos = MyRoPE(config.dim // config.heads)
+        return [
+            TransformerLayer(
+                self_attn=SelfAttention(
+                    SelfAttentionConfig(dim=config.dim, heads=config.heads,
+                                        dim_head=config.dim // config.heads),
+                    pos_encoding=pos,
+                ),
+                ff=SwiGLU(SwiGLUConfig(dim=config.dim, mult=4)),
+                norm_attn=nn.RMSNorm(config.dim),
+                norm_ff=nn.RMSNorm(config.dim),
+            )
+            for _ in range(config.num_layers)
+        ]
+
+    def build_norm(self, config: MyConfig) -> Norm:
+        return nn.RMSNorm(config.dim)
+```

stackformers/presets/__init__.py

@@ -0,0 +1,39 @@
+from stackformers.feedforward.factory import build_ff
+from stackformers.norm.factory import NormConfig, build_norm
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.presets.cross_attender import (
+    CrossAttender,
+    CrossAttenderConfig,
+    plain_cross_attender_config,
+)
+from stackformers.presets.decoder import (
+    TransformerDecoder,
+    TransformerDecoderConfig,
+    plain_decoder_config,
+)
+from stackformers.presets.encoder import (
+    TransformerEncoder,
+    TransformerEncoderConfig,
+    plain_encoder_config,
+    windowed_encoder_config,
+)
+
+__all__ = [
+    # builders
+    "NormConfig",
+    "build_norm",
+    "build_ff",
+    "build_pos_encoding",
+    # preset config factories
+    "plain_encoder_config",
+    "windowed_encoder_config",
+    "plain_decoder_config",
+    "plain_cross_attender_config",
+    # presets
+    "TransformerEncoderConfig",
+    "TransformerEncoder",
+    "TransformerDecoderConfig",
+    "TransformerDecoder",
+    "CrossAttenderConfig",
+    "CrossAttender",
+]

stackformers/presets/cross_attender.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+import torch.nn as nn
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from stackformers.attention.config import CrossAttentionConfig
+from stackformers.attention.cross_attn import CrossAttention
+from stackformers.cross_attender import CrossAttenderLayer, CrossAttenderStack
+from stackformers.feedforward.config import FeedForwardConfig, SwiGLUConfig
+from stackformers.feedforward.factory import build_ff
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import NormConfig, build_norm
+from stackformers.norm.protocols import Norm
+from stackformers.positional.config import NoPosEncodingConfig, PosEncodingConfig
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.positional.protocols import PosEncoding
+from stackformers.sequence import SequenceInput
+
+C = TypeVar("C")
+
+
+class CrossAttenderBase(nn.Module, Generic[C], ABC):
+    """Abstract cross-attender: queries from x attend to context, no self-attention.
+
+    Subclass with any config type C.  Implement build_layers and build_norm;
+    __init__ wires them into a CrossAttenderStack and nothing else.
+    """
+
+    def __init__(self, config: C) -> None:
+        super().__init__()
+        self.config = config
+        self._stack = CrossAttenderStack(
+            layers=self.build_layers(config),
+            final_norm=self.build_norm(config),
+        )
+
+    @abstractmethod
+    def build_layers(self, config: C) -> list[CrossAttenderLayer]: ...
+
+    @abstractmethod
+    def build_norm(self, config: C) -> Norm: ...
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor:
+        return self._stack(x_input, ctx_input)
+
+
+class CrossAttenderConfig(BaseModel):
+    attn: CrossAttentionConfig
+    ff: FeedForwardConfig
+    norm: NormConfig
+    pos_encoding: PosEncodingConfig = NoPosEncodingConfig()
+    num_layers: int = Field(gt=0)
+
+
+def plain_cross_attender_config(
+    dim: int,
+    heads: int,
+    num_layers: int,
+    *,
+    ff_mult: float = 4.0,
+    dropout: float = 0.0,
+) -> CrossAttenderConfig:
+    """Global SDPA cross-attender with RMSNorm and SwiGLU FF, no positional encoding.
+
+    Pass PaddedInput for inference, PackedInput for training — same model.
+    """
+    dim_head = dim // heads
+    return CrossAttenderConfig(
+        attn=CrossAttentionConfig(dim=dim, heads=heads, dim_head=dim_head, dropout=dropout),
+        ff=SwiGLUConfig(dim=dim, mult=ff_mult, dropout=dropout),
+        norm=RMSNormConfig(dim=dim),
+        num_layers=num_layers,
+    )
+
+
+class CrossAttender(CrossAttenderBase[CrossAttenderConfig]):
+    """Concrete cross-attender for CrossAttenderConfig.
+
+    Pass PaddedInput for inference, PackedInput for training.
+    Override build_pos_encoding, build_ff, or build_norm to customise individual
+    collaborators while keeping the rest of the defaults.
+    """
+
+    def build_layers(self, config: CrossAttenderConfig) -> list[CrossAttenderLayer]:
+        pos = self.build_pos_encoding(config)
+        return [
+            CrossAttenderLayer(
+                cross_attn=CrossAttention(config=config.attn, pos_encoding=pos),
+                ff=self.build_ff(config),
+                norm_cross=build_norm(config.norm),
+                norm_ff=build_norm(config.norm),
+            )
+            for _ in range(config.num_layers)
+        ]
+
+    def build_pos_encoding(self, config: CrossAttenderConfig) -> PosEncoding:
+        return build_pos_encoding(config.pos_encoding)
+
+    def build_ff(self, config: CrossAttenderConfig) -> FeedForward:
+        return build_ff(config.ff)
+
+    def build_norm(self, config: CrossAttenderConfig) -> Norm:
+        return build_norm(config.norm)