stackformers 4.2.1__tar.gz → 4.3.0__tar.gz

{stackformers-4.2.1 → stackformers-4.3.0}/CHANGELOG.md

@@ -6,6 +6,37 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
 adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
 MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
 
+## [4.3.0] — 2026-07-17
+
+### Added
+
+- **`MLMWrapper`** (`stackformers/mlm/`) — a masked-token-reconstruction auxiliary loss over any
+  encoder satisfying `EncoderLike`, domain-agnostic about what a token represents. Takes the
+  encoder at call time (`forward(input, encoder)`) rather than owning it, so nothing is
+  registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+- `MLMWrapperConfig`, `MaskingStrategy`/`RandomMasking`, `ReconstructionHead`/`RegressionHead`,
+  and `MLMOutput` (`out`, `mlm_loss`).
+
+### Notes
+
+- **`out` is always the encoder's clean, unmasked output; only `mlm_loss` ever reflects
+  masking.** In training, `forward` runs the encoder twice — once clean (returned as `out`) and
+  once on a separately-masked copy (used only to compute `mlm_loss`) — so the main pipeline sees
+  byte-identical output whether or not this loss is being trained alongside it. In eval, only the
+  clean pass runs and `mlm_loss` is a constant zero, gated on `self.training` the same way
+  `nn.Dropout` and `nn.BatchNorm` already are — so callers can invoke it unconditionally in both
+  modes with no `if training` branch of their own.
+
+- **The reconstruction target is always `input.x.detach()`.** This severs the gradient path from
+  the loss back to whatever produced `input.x`, removing the collapse shortcut a trainable
+  tokenizer would otherwise have (drive every token toward one constant vector to make
+  reconstruction trivial).
+
+- **`RandomMasking` needs no packed-sequence boundary awareness** — each token's masking decision
+  is independent of every other's, so document identity (`cu_seqlens`) never enters it. That only
+  becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle
+  two packed documents.
+
 ## [4.2.1] — 2026-07-14
 
 ### Removed

{stackformers-4.2.1 → stackformers-4.3.0}/CLAUDE.md

@@ -25,6 +25,7 @@ Behavior comes from injected dependencies, not constructor flags.
 | dh     | dim per head |
 | w      | window size |
 | nt     | total tokens in packed sequence |
+| m      | number of masked tokens (gathered across batch/pack) |
 
 ## File organisation
 

{stackformers-4.2.1 → stackformers-4.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 4.2.1
+Version: 4.3.0
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-4.2.1 → stackformers-4.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "4.2.1"
+version = "4.3.0"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-4.2.1 → stackformers-4.3.0}/stackformers/__init__.py

@@ -17,6 +17,11 @@ from stackformers.feedforward.factory import build_ff
 from stackformers.feedforward.protocols import FeedForward
 from stackformers.feedforward.swiglu import SwiGLU
 from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
 from stackformers.norm.config import LayerNormConfig, RMSNormConfig
 from stackformers.norm.factory import NormConfig, build_norm
 from stackformers.norm.protocols import Norm
@@ -85,6 +90,9 @@ __all__ = [
     "CrossAttn",
     "FeedForward",
     "Norm",
+    "EncoderLike",
+    "MaskingStrategy",
+    "ReconstructionHead",
     # configs — attention
     "SelfAttentionConfig",
     "CrossAttentionConfig",
@@ -133,4 +141,10 @@ __all__ = [
     "CrossAttenderConfig",
     "CrossAttender",
     "plain_cross_attender_config",
+    # mlm
+    "MLMWrapperConfig",
+    "RandomMasking",
+    "RegressionHead",
+    "MLMWrapper",
+    "MLMOutput",
 ]

stackformers-4.3.0/stackformers/mlm/README.md

@@ -0,0 +1,54 @@
+# mlm
+
+A masked-token-reconstruction auxiliary loss, wrapped around any encoder that satisfies `EncoderLike`. Domain-agnostic — it has no idea what a token represents.
+
+## Usage
+
+```python
+import torch
+from stackformers import (
+    MLMWrapper,
+    MLMWrapperConfig,
+    TransformerEncoder,
+    make_padded_input,
+    plain_encoder_config,
+)
+
+encoder = TransformerEncoder(plain_encoder_config(dim=512, heads=8, num_layers=6))
+mlm_wrapper = MLMWrapper(MLMWrapperConfig(dim=512, mask_ratio=0.15))
+
+x = torch.randn(2, 128, 512)
+mask = torch.ones(2, 128, dtype=torch.bool)
+input = make_padded_input(x, mask)
+
+res = mlm_wrapper(input, encoder)
+
+task_loss = my_task_head(res.out)  # res.out is always the clean, unmasked encoder output
+loss = task_loss + 0.1 * res.mlm_loss  # zero automatically in eval — nothing to branch on
+```
+
+One call site covers both modes: in training, `mlm_wrapper` runs the encoder twice internally — once clean (returned as `res.out`) and once on a separately-masked copy (used only to compute `res.mlm_loss`) — so masking never reaches whatever consumes `res.out`. In eval, it runs the encoder once and `res.mlm_loss` is a constant zero. Pretraining with no main task yet — just drop `task_loss` and train on `res.mlm_loss` alone.
+
+`mlm_wrapper` and `encoder` are separate modules — `encoder` is passed in at call time, never stored, so `mlm_wrapper.parameters()` holds only `mask_token`, `masking_strategy`, and `head`. Build one optimizer over both: `itertools.chain(encoder.parameters(), mlm_wrapper.parameters())`. And since they're separate modules, `.train()`/`.eval()` on one does not propagate to the other — call both, or register both as submodules of a shared parent so one call does.
+
+## Design
+
+**`out` is always clean; only `mlm_loss` ever reflects masking.** `forward(input, encoder)` runs `encoder(input)` unconditionally for `out`. In training it additionally runs `encoder` a second time on a masked copy, purely to produce `mlm_loss` — that corrupted copy is never returned. This means the main pipeline sees byte-identical output whether or not the aux loss is being trained alongside it, and callers never need to special-case which pass produced `out`.
+
+**Training/eval is the one branch `forward()` is allowed, because it isn't tensor control flow.** `self.training` gates whether the second (masked) pass runs — the same mechanism `nn.Dropout` and `nn.BatchNorm` already use. It lets callers invoke `mlm_wrapper(input, encoder)` unconditionally in both modes, instead of an external `if training: ...` at every call site.
+
+**A wrapper over an encoder supplied at call time, not one it owns.** `MLMWrapper` never stores an encoder — `forward` takes one as an argument, so nothing is registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+
+**The mask token is owned state, not an injected collaborator.** It's the wrapper's own learned parameter, analogous to a `nn.Linear`'s weight — not swappable behavior, so it isn't a `Protocol`.
+
+**The reconstruction target is always `input.x.detach()`, unconditionally.** Not a config flag: computing a detached target is the only behavior the masked pass ever has, so there is nothing to branch on. This severs the gradient path from the reconstruction loss back to whatever produced `input.x`, which removes the standard collapse shortcut for self-supervised regression targets (drive every token toward one constant vector, and reconstruction becomes trivial). See `tests/mlm/test_wrapper.py` for the gradient test that verifies this directly.
+
+**No layout dispatch in the wrapper.** `MaskingStrategy` returns a boolean tensor shaped like `input.x`'s leading dims — `(b, n)` for `PaddedInput`, `(nt,)` for `PackedInput`. `torch.where` and boolean advanced indexing both broadcast/gather correctly against either shape, so `MLMWrapper.forward` never matches on the sequence variant; only `RandomMasking` does, internally.
+
+**Random masking needs no packed-sequence boundary awareness.** Each token's masking decision is independent of every other token's, so document identity (`cu_seqlens`) never enters the decision. Boundary-awareness only becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle two packed documents — add that logic when that strategy is added, not before.
+
+## Extending
+
+To add a masking strategy (e.g. blockwise), write an `nn.Module` satisfying `MaskingStrategy` and pass it as `MLMWrapper(..., masking_strategy=...)`.
+
+To add a reconstruction target type (e.g. a discretized/codebook target for BEiT-style cross-entropy, or a contrastive target), write an `nn.Module` satisfying `ReconstructionHead` and pass it as `MLMWrapper(..., head=...)`. Prediction and loss are scored together by one call — they're a coupled choice (a cross-entropy loss over a raw regression output isn't meaningful), not two independently pluggable axes.

stackformers-4.3.0/stackformers/mlm/config.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class MLMWrapperConfig(BaseModel):
+    """Config for MLMWrapper: token dimension and corruption ratio."""
+
+    dim: int = Field(gt=0)
+    mask_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)

stackformers-4.3.0/stackformers/mlm/head.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import torch.nn as nn
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
+
+
+class RegressionHead(nn.Module):
+    """Reconstruct masked tokens via a linear projection, scored with MSE.
+
+    The projection gives the encoder a dedicated place to specialise for
+    reconstruction, so the shared representation isn't forced to double as the
+    literal token vector at every layer.
+    """
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]:
+        prediction = self.proj(encoder_output_at_masked)
+        return F.mse_loss(prediction, target_at_masked)

stackformers-4.3.0/stackformers/mlm/masking.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from jaxtyping import Bool
+from torch import Tensor
+
+from stackformers.sequence import PackedInput, PaddedInput, SequenceInput
+
+
+class RandomMasking(nn.Module):
+    """Mask each valid token independently with probability mask_ratio.
+
+    Every token's masking decision is drawn independently, so packed sequences need
+    no document-boundary awareness here — that only becomes necessary for a
+    contiguous-span (blockwise) strategy, which could otherwise straddle two packed
+    documents.
+    """
+
+    def __init__(self, mask_ratio: float = 0.15) -> None:
+        super().__init__()
+        self.mask_ratio = mask_ratio
+
+    def forward(self, input: SequenceInput) -> Bool[Tensor, "*batch"]:
+        match input:
+            case PaddedInput(mask=mask):
+                candidate = torch.rand(mask.shape, device=mask.device) < self.mask_ratio
+                return candidate & mask
+            case PackedInput(x=x):
+                return torch.rand(x.shape[0], device=x.device) < self.mask_ratio

stackformers-4.3.0/stackformers/mlm/protocols.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+from stackformers.sequence import SequenceInput
+
+
+@runtime_checkable
+class EncoderLike(Protocol):
+    """Any encoder mapping a sequence to per-token embeddings of the same layout.
+
+    Implementations: Encoder, or any TransformerEncoderBase subclass.
+    """
+
+    def __call__(self, input: SequenceInput) -> Tensor: ...
+
+
+@runtime_checkable
+class MaskingStrategy(Protocol):
+    """Select which token positions to corrupt for reconstruction.
+
+    Returns True at positions to mask, shaped like the leading (non-feature) dims of
+    input.x — (b, n) for PaddedInput, (nt,) for PackedInput. Implementations dispatch
+    on the SequenceInput variant internally, so callers never need to know the layout.
+    Implementation: RandomMasking.
+    """
+
+    def __call__(self, input: SequenceInput) -> Bool[Tensor, "*batch"]: ...
+
+
+@runtime_checkable
+class ReconstructionHead(Protocol):
+    """Predict masked tokens from encoder output and score against the clean target.
+
+    Both arguments are already gathered down to just the masked positions (m = number
+    of masked tokens across the batch or pack). Returns a scalar loss.
+    Implementation: RegressionHead.
+    """
+
+    def __call__(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]: ...

stackformers-4.3.0/stackformers/mlm/wrapper.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.sequence import SequenceInput
+
+
+class MLMOutput(NamedTuple):
+    out: Float[Tensor, "*batch d"]
+    mlm_loss: Float[Tensor, ""]
+
+
+class MLMWrapper(nn.Module):
+    """Masked-token-reconstruction auxiliary loss over an encoder supplied at call time.
+
+    Does not store or own an encoder — none is passed to __init__, and forward() takes
+    one as an argument instead of holding a reference across calls. This module's own
+    parameter tree (mask_token, masking_strategy, head) is exactly what it owns; the
+    encoder stays wherever the caller already keeps it, with no duplicate registration
+    under this wrapper.
+
+    `out` is always the encoder's clean, unmasked output — masking is invisible to
+    whatever consumes it, so the main pipeline behaves identically whether or not this
+    aux loss is being trained alongside it. Only `mlm_loss` ever reflects that masking
+    happened. In training mode that costs a second encoder forward pass on a
+    separately-masked copy of input; the two passes share weights but never interfere,
+    since the corrupted copy built for the loss is never returned as `out`. In eval mode
+    only the one clean pass runs and `mlm_loss` reports a constant zero, so callers can
+    invoke this unconditionally in both modes without an if-training branch of their
+    own — the same role self.training already plays in nn.Dropout or nn.BatchNorm.
+
+    The reconstruction target is always the encoder's own clean input embedding,
+    detached — so this loss trains the encoder and head, never whatever produced
+    input.x. That removes the representation-collapse shortcut a trainable tokenizer
+    would otherwise have available (drive every token toward a constant vector to make
+    reconstruction trivial).
+    """
+
+    def __init__(
+        self,
+        config: MLMWrapperConfig,
+        masking_strategy: MaskingStrategy | None = None,
+        head: ReconstructionHead | None = None,
+    ) -> None:
+        super().__init__()
+        self.mask_token = nn.Parameter(torch.empty(config.dim))
+        nn.init.trunc_normal_(self.mask_token, std=0.02, a=-0.04, b=0.04)
+        self.masking_strategy = (
+            masking_strategy if masking_strategy is not None else RandomMasking(config.mask_ratio)
+        )
+        self.head = head if head is not None else RegressionHead(config.dim)
+
+    def _masked_loss(self, input: SequenceInput, encoder: EncoderLike) -> Tensor:
+        should_mask = self.masking_strategy(input)
+        target = input.x.detach()
+        # should_mask's shape always matches input.x's leading dims, so this select and
+        # the boolean indexing below both work for PaddedInput and PackedInput alike.
+        corrupted_x = torch.where(should_mask.unsqueeze(-1), self.mask_token, input.x)
+        masked_output = encoder(input._replace(x=corrupted_x))
+        return self.head(masked_output[should_mask], target[should_mask])
+
+    def forward(self, input: SequenceInput, encoder: EncoderLike) -> MLMOutput:
+        clean_output = encoder(input)
+        if self.training:
+            mlm_loss = self._masked_loss(input, encoder)
+            return MLMOutput(out=clean_output, mlm_loss=mlm_loss)
+        return MLMOutput(out=clean_output, mlm_loss=clean_output.new_zeros(()))

stackformers-4.3.0/tests/mlm/test_head.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from stackformers.mlm.head import RegressionHead
+
+M, D = 12, 16
+
+
+def test_regression_head_output_is_scalar(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    head = RegressionHead(dim=D).to(device=device, dtype=dtype)
+    prediction_in = torch.randn(M, D, device=device, dtype=dtype)
+    target = torch.randn(M, D, device=device, dtype=dtype)
+    loss = head(prediction_in, target)
+    assert loss.shape == ()
+
+
+def test_regression_head_zero_when_projection_matches_target(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    nn.init.zeros_(head.proj.weight)
+    nn.init.zeros_(head.proj.bias)
+    target = torch.zeros(M, D, device=device)
+    prediction_in = torch.randn(M, D, device=device)  # proj(x) == 0 regardless, weight/bias are 0
+    loss = head(prediction_in, target)
+    assert torch.allclose(loss, torch.zeros_like(loss))
+
+
+def test_regression_head_gradients_flow(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    prediction_in = torch.randn(M, D, device=device, requires_grad=True)
+    target = torch.randn(M, D, device=device)
+    loss = head(prediction_in, target)
+    loss.backward()
+    assert prediction_in.grad is not None
+    assert head.proj.weight.grad is not None

stackformers-4.3.0/tests/mlm/test_masking.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import torch
+
+from stackformers.mlm.masking import RandomMasking
+from stackformers.sequence import make_packed_input, make_padded_input
+
+B, N, D = 4, 64, 16
+NT = 20
+
+
+def test_random_masking_padded_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == mask.shape
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_packed_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 12, 20], dtype=torch.int32, device=device)
+    input = make_packed_input(x, cu, max_seqlen=12)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == (NT,)
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_never_masks_padding(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    mask[:, N // 2 :] = False  # second half is padding
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.9)(input)
+    assert not (should_mask & ~mask).any()
+
+
+def test_random_masking_respects_ratio_roughly(
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """Statistical check over a large sample — not exact, but far enough from any other ratio."""
+    device, dtype = device_dtype
+    torch.manual_seed(0)
+    x = torch.randn(1, 10_000, D, device=device, dtype=dtype)
+    mask = torch.ones(1, 10_000, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.3)(input)
+    fraction = should_mask.float().mean().item()
+    assert 0.25 < fraction < 0.35

stackformers-4.3.0/tests/mlm/test_wrapper.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from stackformers.attention.config import SelfAttentionConfig
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.encoder import Encoder
+from stackformers.feedforward.config import SwiGLUConfig
+from stackformers.feedforward.swiglu import SwiGLU
+from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import build_norm
+from stackformers.positional.none import NoPosEncoding
+from stackformers.sequence import (
+    PackedInput,
+    PaddedInput,
+    SequenceInput,
+    make_packed_input,
+    make_padded_input,
+)
+
+B, N, D, H, DH = 2, 16, 16, 4, 4
+NUM_LAYERS = 2
+NT = 10  # two packed seqs: 6 + 4
+
+
+class AllMasking(nn.Module):
+    """Test double: marks every valid position for masking."""
+
+    def forward(self, input: SequenceInput) -> torch.Tensor:
+        match input:
+            case PaddedInput(mask=mask):
+                return mask.clone()
+            case PackedInput(x=x):
+                return torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
+
+
+def _build_encoder(device: torch.device, dtype: torch.dtype) -> Encoder:
+    attn_cfg = SelfAttentionConfig(dim=D, heads=H, dim_head=DH)
+    ff_cfg = SwiGLUConfig(dim=D)
+    norm_cfg = RMSNormConfig(dim=D)
+    layers = [
+        TransformerLayer(
+            self_attn=SelfAttention(attn_cfg, NoPosEncoding()),
+            ff=SwiGLU(ff_cfg),
+            norm_attn=build_norm(norm_cfg),
+            norm_ff=build_norm(norm_cfg),
+        )
+        for _ in range(NUM_LAYERS)
+    ]
+    return Encoder(layers=layers, final_norm=build_norm(norm_cfg)).to(device=device, dtype=dtype)
+
+
+@pytest.fixture
+def config() -> MLMWrapperConfig:
+    return MLMWrapperConfig(dim=D, mask_ratio=0.5)
+
+
+@pytest.fixture
+def x_pad(device_dtype: tuple[torch.device, torch.dtype]) -> PaddedInput:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    return make_padded_input(x, mask)
+
+
+@pytest.fixture
+def x_packed(device_dtype: tuple[torch.device, torch.dtype]) -> PackedInput:
+    device, dtype = device_dtype
+    if not device.type == "cuda" or dtype not in (torch.float16, torch.bfloat16):
+        pytest.skip("packed attention requires CUDA with float16 or bfloat16")
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 6, 10], dtype=torch.int32, device=device)
+    return make_packed_input(x, cu, max_seqlen=6)
+
+
+def test_wrapper_padded_output_shapes(
+    config: MLMWrapperConfig,
+    x_pad: PaddedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_pad, encoder)
+    assert isinstance(res, MLMOutput)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (B, N, D)
+
+
+def test_wrapper_packed_output_shapes(
+    config: MLMWrapperConfig,
+    x_packed: PackedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_packed, encoder)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (NT, D)
+
+
+def test_wrapper_mask_token_receives_gradient(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert wrapper.mask_token.grad is not None
+
+
+def test_wrapper_does_not_own_encoder(device: torch.device) -> None:
+    """MLMWrapper takes the encoder at call time, not at construction — its own
+    parameter tree must never include the encoder's weights, since the encoder is
+    owned and registered wherever the caller already keeps it.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    encoder_param_ids = {id(p) for p in encoder.parameters()}
+    wrapper_param_ids = {id(p) for p in wrapper.parameters()}
+    assert encoder_param_ids.isdisjoint(wrapper_param_ids)
+
+
+def test_wrapper_detaches_target_from_upstream_input(device: torch.device) -> None:
+    """Collapse guard (design doc §5): a trainable tokenizer upstream of x must get no
+    gradient from this loss, or the optimizer can collapse every token to one constant
+    vector to make reconstruction trivial.
+
+    Force every position masked, so torch.where's backward zeroes x's contribution
+    through the corrupted-x branch entirely (its condition is True everywhere, so the
+    "x" branch of the select is never taken). The only remaining path from loss back to
+    x would be through the target — and that's detached — so x's accumulated gradient
+    must be exactly zero.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device, requires_grad=True)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert x.grad is not None
+    assert torch.all(x.grad == 0)
+
+
+def test_wrapper_accepts_custom_masking_strategy(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    assert res.mlm_loss.shape == ()
+
+
+def test_wrapper_training_out_is_always_clean(device: torch.device) -> None:
+    """out must never carry masking, even in training — only mlm_loss does. This keeps
+    the main pipeline byte-identical whether or not the MLM aux loss trains alongside
+    it, so callers never need to special-case which pass produced `out`. AllMasking
+    maximizes the contrast: if masking leaked into `out`, every token would be replaced
+    by mask_token, making it maximally different from a clean encoder(input) call.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+
+
+def test_wrapper_eval_mode_passes_through_encoder_unmasked(device: torch.device) -> None:
+    """In eval mode, forward() must skip masking entirely: out matches a clean encoder
+    call exactly, and mlm_loss reports zero — so callers can invoke the wrapper
+    unconditionally in both modes without an if-training branch of their own.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    wrapper.eval()
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+    assert torch.equal(res.mlm_loss, torch.zeros_like(res.mlm_loss))

{stackformers-4.2.1 → stackformers-4.3.0}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "4.2.1"
+version = "4.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },