Fix transformers 4.x/5.x compat, implement output_hidden_states/attentions and out_layers, fix hierarchical predictor input, add video processor

__init__.py

@@ -1,15 +1,25 @@
+"""V-JEPA 2.1 — HuggingFace port."""
+
 from .configuration_vjepa21 import VJEPA21Config
 from .modeling_vjepa21 import (
     VJEPA21ForVideoClassification,
     VJEPA21Model,
     VJEPA21PreTrainedModel,
 )
-from .video_processing_vjepa21 import VJEPA21VideoProcessor
 
 __all__ = [
     "VJEPA21Config",
     "VJEPA21Model",
     "VJEPA21PreTrainedModel",
     "VJEPA21ForVideoClassification",
-    "VJEPA21VideoProcessor",
 ]
+
+# `BaseVideoProcessor` needs torchvision. Import it lazily so that a runtime
+# without torchvision can still load the model; `AutoVideoProcessor` resolves the
+# class through `auto_map` and does not go through this file.
+try:  # pragma: no cover - depends on the environment
+    from .video_processing_vjepa21 import VJEPA21VideoProcessor
+
+    __all__.append("VJEPA21VideoProcessor")
+except ImportError:  # pragma: no cover
+    pass

config.json

@@ -28,8 +28,8 @@
   "n_registers": 0,
   "num_attention_heads": 12,
   "num_hidden_layers": 12,
-  "num_pooler_layers": 3,
   "num_pooler_heads": 16,
+  "num_pooler_layers": 3,
   "patch_size": 16,
   "pred_hidden_size": 384,
   "pred_mlp_ratio": 4.0,

configuration_vjepa21.py

@@ -5,26 +5,52 @@ from typing import Optional
 from transformers import PretrainedConfig
 
 
+# Encoder hierarchical layers, from `app/vjepa_2_1/models/vision_transformer.py`
+# in `facebookresearch/vjepa2`.
+_ENCODER_LAYER_MAP = {
+    4: [0, 1, 2, 3],
+    8: [1, 3, 5, 7],
+    12: [2, 5, 8, 11],
+    20: [4, 9, 14, 19],
+    24: [5, 11, 17, 23],
+    40: [9, 19, 29, 39],
+    48: [11, 23, 37, 47],
+}
+
+# Predictor hierarchical layers, from `app/vjepa_2_1/models/predictor.py`.
+# NOTE: this is *not* the same table as the encoder's. At depth 24 the reference
+# predictor uses [4, 11, 17, 23] while the encoder uses [5, 11, 17, 23], and the
+# predictor table has no entry for depth 48. Sharing one table happens to work
+# today because only `len(...)` is consumed, but it would silently produce wrong
+# indices as soon as multi-level predictor outputs are exposed.
+_PREDICTOR_LAYER_MAP = {
+    4: [0, 1, 2, 3],
+    8: [1, 3, 5, 7],
+    12: [2, 5, 8, 11],
+    20: [4, 9, 14, 19],
+    24: [4, 11, 17, 23],
+    40: [9, 19, 29, 39],
+}
+
+
 def _get_hierarchical_layers(depth: int) -> list[int]:
-    """Get hierarchical layer indices based on model depth.
+    """Encoder hierarchical layer indices for a given depth."""
+    if depth not in _ENCODER_LAYER_MAP:
+        raise ValueError(
+            f"Unsupported encoder depth {depth}. Supported depths: "
+            f"{list(_ENCODER_LAYER_MAP.keys())}"
+        )
+    return _ENCODER_LAYER_MAP[depth]
 
-    Mirrors `app/vjepa_2_1/models/vision_transformer.py` in the reference
-    implementation.
-    """
-    _LAYER_MAP = {
-        4: [0, 1, 2, 3],
-        8: [1, 3, 5, 7],
-        12: [2, 5, 8, 11],
-        20: [4, 9, 14, 19],
-        24: [5, 11, 17, 23],
-        40: [9, 19, 29, 39],
-        48: [11, 23, 37, 47],
-    }
-    if depth not in _LAYER_MAP:
+
+def _get_predictor_hierarchical_layers(depth: int) -> list[int]:
+    """Predictor hierarchical layer indices for a given depth."""
+    if depth not in _PREDICTOR_LAYER_MAP:
         raise ValueError(
-            f"Unsupported depth {depth}. Supported depths: {list(_LAYER_MAP.keys())}"
+            f"Unsupported predictor depth {depth}. Supported depths: "
+            f"{list(_PREDICTOR_LAYER_MAP.keys())}"
         )
-    return _LAYER_MAP[depth]
+    return _PREDICTOR_LAYER_MAP[depth]
 
 
 class VJEPA21Config(PretrainedConfig):
@@ -86,10 +112,15 @@ class VJEPA21Config(PretrainedConfig):
         has_cls_first (`bool`, defaults to False):
             Whether the sequence starts with a CLS token.
         num_pooler_layers (`int`, defaults to 3):
-            Number of self-attention layers in the attentive pooler.
-        num_pooler_heads (`int` or `None`, defaults to None):
-            Number of attention heads in the attentive pooler. Defaults to
-            `num_attention_heads`. The reference frozen-evaluation probes use 16.
+            Number of self-attention layers in the attentive pooler. Together with
+            the cross-attention layer this reproduces `AttentivePooler(depth=4)`,
+            which is `num_probe_blocks: 4` in the reference evaluation configs.
+        num_pooler_heads (`int` or `None`, defaults to 16):
+            Number of attention heads in the attentive pooler. 16 is the value used
+            by every frozen-probe config under `configs/eval_2_1/` in the reference
+            repository (`classifier.num_heads: 16`), for all four model sizes, so it
+            is the default here rather than `num_attention_heads`. The pooler is
+            always trained from scratch, so this only affects the probe you train.
         pred_hidden_size (`int`, defaults to 384):
             Predictor embedding dimension.
         pred_num_attention_heads (`int`, defaults to 12):
@@ -138,7 +169,7 @@ class VJEPA21Config(PretrainedConfig):
         has_cls_first: bool = False,
         # Pooler
         num_pooler_layers: int = 3,
-        num_pooler_heads: Optional[int] = None,
+        num_pooler_heads: Optional[int] = 16,
         # Predictor
         pred_hidden_size: int = 384,
         pred_num_attention_heads: int = 12,
@@ -179,9 +210,7 @@ class VJEPA21Config(PretrainedConfig):
 
         # Pooler
         self.num_pooler_layers = num_pooler_layers
-        self.num_pooler_heads = (
-            num_pooler_heads if num_pooler_heads is not None else num_attention_heads
-        )
+        self.num_pooler_heads = num_pooler_heads if num_pooler_heads is not None else 16
 
         # Predictor
         self.pred_hidden_size = pred_hidden_size
@@ -197,11 +226,18 @@ class VJEPA21Config(PretrainedConfig):
 
     def _validate(self) -> None:
         n_levels = len(_get_hierarchical_layers(self.num_hidden_layers))
+        n_pred_levels = len(_get_predictor_hierarchical_layers(self.pred_num_hidden_layers))
         if not 1 <= self.n_output_distillation <= n_levels:
             raise ValueError(
                 f"n_output_distillation must be in [1, {n_levels}] for a model with "
                 f"{self.num_hidden_layers} layers, got {self.n_output_distillation}."
             )
+        if self.n_output_distillation > n_pred_levels:
+            raise ValueError(
+                f"n_output_distillation ({self.n_output_distillation}) exceeds the "
+                f"{n_pred_levels} hierarchical levels available in a predictor with "
+                f"{self.pred_num_hidden_layers} layers."
+            )
         if self.hidden_size % self.num_attention_heads != 0:
             raise ValueError(
                 f"hidden_size ({self.hidden_size}) must be divisible by "
@@ -217,22 +253,38 @@ class VJEPA21Config(PretrainedConfig):
                 f"pred_hidden_size ({self.pred_hidden_size}) must be divisible by "
                 f"pred_num_attention_heads ({self.pred_num_attention_heads})."
             )
+        if self.pred_teacher_embed_dim is not None:
+            if self.pred_teacher_embed_dim % self.n_output_distillation != 0:
+                raise ValueError(
+                    f"pred_teacher_embed_dim ({self.pred_teacher_embed_dim}) must be "
+                    f"divisible by n_output_distillation ({self.n_output_distillation})."
+                )
+        if self.tubelet_size < 1:
+            raise ValueError(f"tubelet_size must be >= 1, got {self.tubelet_size}.")
+        if self.pred_num_mask_tokens < 1:
+            raise ValueError(
+                f"pred_num_mask_tokens must be >= 1, got {self.pred_num_mask_tokens}."
+            )
 
     @property
     def encoder_hierarchical_layers(self) -> list[int]:
-        """Layer indices for hierarchical output collection in the encoder."""
+        """Layer indices at which the encoder carries a per-level LayerNorm."""
         return _get_hierarchical_layers(self.num_hidden_layers)
 
     @property
     def encoder_distillation_layers(self) -> list[int]:
-        """Layer indices used for distillation output in the encoder."""
+        """Encoder layer indices contributing to the hierarchical output."""
         all_layers = _get_hierarchical_layers(self.num_hidden_layers)
         return all_layers[-self.n_output_distillation :]
 
     @property
     def predictor_hierarchical_layers(self) -> list[int]:
-        """Layer indices for hierarchical output in the predictor."""
-        all_layers = _get_hierarchical_layers(self.pred_num_hidden_layers)
+        """Predictor layer indices for hierarchical output.
+
+        Uses the predictor's own depth table, which differs from the encoder's at
+        depth 24 (`[4, 11, 17, 23]` vs `[5, 11, 17, 23]`).
+        """
+        all_layers = _get_predictor_hierarchical_layers(self.pred_num_hidden_layers)
         return all_layers[-self.n_output_distillation :]
 
     @property

modeling_vjepa21.py

@@ -10,6 +10,14 @@ Key differences from V-JEPA 2:
 - Dense predictor: hierarchical input fusion + context token prediction
 
 Compatible with transformers >= 4.50 (both the 4.x and 5.x attention APIs).
+
+Note on initialisation: the reference `VisionTransformer` and
+`VisionTransformerPredictor` call `_rescale_blocks()` after `_init_weights`,
+dividing `attn.proj.weight` and `mlp.fc2.weight` of layer *i* by `sqrt(2*(i+1))`.
+That is deliberately **not** reproduced here: it only affects randomly
+initialised models, and applying it inside `_init_weights` would risk touching
+weights that `from_pretrained` has already loaded. Every published checkpoint
+loads its weights, so the two agree.
 """
 
 from __future__ import annotations
@@ -30,6 +38,22 @@ from transformers.utils import ModelOutput, logging
 from .configuration_vjepa21 import VJEPA21Config
 
 logger = logging.get_logger(__name__)
+_warn_once = getattr(logger, "warning_once", logger.warning)
+
+
+# Keyword arguments the model forward understands. Anything else is reported
+# once and dropped, instead of being swallowed by `**kwargs`: a typo such as
+# `out_layer=[11]` used to be a silent no-op.
+_ENCODER_FORWARD_KWARGS = frozenset(
+    {"masks", "out_layers", "return_hierarchical", "output_attentions", "output_hidden_states"}
+)
+_MODEL_FORWARD_KWARGS = _ENCODER_FORWARD_KWARGS | frozenset(
+    {"context_mask", "target_mask", "skip_predictor", "mask_index"}
+)
+# Injected by the Trainer or by generic HF plumbing; harmless and not worth a warning.
+_SILENTLY_IGNORED_KWARGS = frozenset(
+    {"num_items_in_batch", "return_dict", "return_loss", "interpolate_pos_encoding"}
+)
 
 
 # ---------------------------------------------------------------------------
@@ -48,8 +72,9 @@ class VJEPA21EncoderOutput(ModelOutput):
             channel axis, `(B, N, n_output_distillation * hidden_size)`.
             Only returned when `return_hierarchical=True`.
         multilevel_hidden_states: Tuple of per-level normalised features, one
-            entry per requested `out_layers` index. Mirrors the `out_layers`
-            argument of the reference implementation.
+            entry per requested `out_layers` index, in the order the layers occur
+            in the network (not in the order they were requested). Mirrors the
+            `out_layers` argument of the reference implementation.
         hidden_states: Raw (un-normalised) outputs of the embedding layer and of
             every transformer layer, when `output_hidden_states=True`.
         attentions: Attention probabilities of every layer, when
@@ -83,7 +108,14 @@ class VJEPA21PredictorOutput(ModelOutput):
 
 @dataclass
 class VJEPA21ModelOutput(ModelOutput):
-    """Full model output combining encoder and predictor."""
+    """Full model output combining encoder and predictor.
+
+    `masked_hidden_state` is the tensor the predictor actually consumed, gathered
+    at `context_mask`. Its channel width therefore depends on the checkpoint:
+    `hidden_size` when `n_output_distillation == 1` (ViT-B, ViT-L) and
+    `n_output_distillation * hidden_size` when the predictor fuses several levels
+    (ViT-g, ViT-G).
+    """
 
     last_hidden_state: Optional[torch.FloatTensor] = None
     hierarchical_hidden_state: Optional[torch.FloatTensor] = None
@@ -117,6 +149,15 @@ def apply_masks(tensor: torch.Tensor, masks: list[torch.Tensor]) -> torch.Tensor
     return torch.cat(parts, dim=0)
 
 
+def _as_mask_list(masks) -> Optional[list[torch.Tensor]]:
+    """Accept a single tensor or a list of tensors, as the reference does."""
+    if masks is None:
+        return None
+    if isinstance(masks, torch.Tensor):
+        return [masks]
+    return list(masks)
+
+
 def normalize_video_layout(pixel_values_videos: torch.Tensor, in_chans: int = 3) -> torch.Tensor:
     """Bring any accepted video layout to channels-first `(B, C, T, H, W)`.
 
@@ -127,7 +168,9 @@ def normalize_video_layout(pixel_values_videos: torch.Tensor, in_chans: int = 3)
         - `(B, C, H, W)`    a batch of single images, promoted to `T = 1`
 
     The channel axis is identified by matching `config.in_chans`. When more than
-    one axis matches (e.g. a 3-frame clip), the channels-first reading wins.
+    one axis matches — a 3-frame clip is the realistic case — the channels-first
+    reading wins and a warning is emitted, because the alternative is a silent
+    transposition.
     """
     x = pixel_values_videos
     if x.ndim == 4:
@@ -144,6 +187,16 @@ def normalize_video_layout(pixel_values_videos: torch.Tensor, in_chans: int = 3)
             f"{x.ndim} dimensions with shape {tuple(x.shape)}."
         )
 
+    candidates = [axis for axis in (1, 2, 4) if x.shape[axis] == in_chans]
+    if len(candidates) > 1:
+        _warn_once(
+            "ambiguous video layout: "
+            f"pixel_values_videos of shape {tuple(x.shape)} has {len(candidates)} candidate "
+            f"channel axes {candidates} of size in_chans={in_chans}; "
+            "reading it as (B, C, T, H, W). "
+            "Pass an unambiguous layout if that is not what you meant."
+        )
+
     if x.shape[1] == in_chans:  # (B, C, T, H, W)
         return x
     if x.shape[2] == in_chans:  # (B, T, C, H, W)
@@ -176,6 +229,9 @@ class VJEPA21DropPath(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return drop_path(x, self.p, self.training)
 
+    def extra_repr(self) -> str:
+        return f"p={self.p}"
+
 
 def rotate_queries_or_keys(
     x: torch.Tensor,
@@ -192,6 +248,9 @@ def rotate_queries_or_keys(
         has_cls_first: Whether first token is CLS (not rotated).
     """
     B, num_heads, N, D = x.size()
+    if D % 2 != 0:
+        raise ValueError(f"RoPE requires an even dimension per slice, got {D}.")
+
     n_cls = 1 if has_cls_first else 0
     start_ctx = n_cls
     end_ctx = N - n_registers
@@ -634,12 +693,31 @@ class VJEPA21Encoder(nn.Module):
     def forward(
         self,
         pixel_values_videos: torch.Tensor,
+        masks: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
         return_hierarchical: bool = False,
         out_layers: Optional[list[int]] = None,
         output_attentions: bool = False,
         output_hidden_states: bool = False,
         **kwargs,
     ) -> VJEPA21EncoderOutput:
+        """
+        Args:
+            masks: Optional list of `(B, K)` index tensors. When given, the patch
+                tokens are gathered at those indices *before* the transformer
+                layers, so attention only ever sees the context tokens and RoPE
+                receives their true positions. This is the JEPA training-time
+                forward (`z = encoder(clips, masks_enc)` in the reference); the
+                default `None` runs the full sequence, which is what feature
+                extraction wants.
+        """
+        unexpected = set(kwargs) - _SILENTLY_IGNORED_KWARGS
+        if unexpected:
+            _warn_once(
+                f"VJEPA21Encoder.forward received unexpected keyword arguments "
+                f"{sorted(unexpected)}; they are ignored. Accepted arguments: "
+                f"{sorted(_ENCODER_FORWARD_KWARGS)}."
+            )
+
         pixel_values_videos = normalize_video_layout(
             pixel_values_videos, self.config.in_chans
         )
@@ -663,11 +741,31 @@ class VJEPA21Encoder(nn.Module):
                     f"hierarchical layers of this model. Valid indices: {self._hier_layers}."
                 )
 
+        # Masked (JEPA) forward: drop tokens before the layers and carry their
+        # original indices as RoPE positions, as the reference does.
+        masks = _as_mask_list(masks)
+        position_mask = None
+        if masks is not None:
+            n_tokens = embeddings.shape[1]
+            for m in masks:
+                if m.dim() != 2 or m.shape[0] != embeddings.shape[0]:
+                    raise ValueError(
+                        f"each mask must be a (B, K) index tensor with B={embeddings.shape[0]}, "
+                        f"got {tuple(m.shape)}."
+                    )
+                if int(m.max()) >= n_tokens:
+                    raise ValueError(
+                        f"mask index {int(m.max())} is out of range for a sequence of "
+                        f"{n_tokens} tokens."
+                    )
+            embeddings = apply_masks(embeddings, masks)
+            position_mask = torch.cat([m.to(embeddings.device) for m in masks], dim=0)
+
         hidden_states = embeddings
         hier_outputs: list[torch.Tensor] = []
         multilevel_outputs: list[torch.Tensor] = []
-        all_hidden_states: tuple[torch.Tensor, ...] = () if output_hidden_states else None
-        all_attentions: tuple[torch.Tensor, ...] = () if output_attentions else None
+        all_hidden_states: Optional[tuple[torch.Tensor, ...]] = () if output_hidden_states else None
+        all_attentions: Optional[tuple[torch.Tensor, ...]] = () if output_attentions else None
 
         for i, layer_module in enumerate(self.layer):
             if output_hidden_states:
@@ -677,7 +775,7 @@ class VJEPA21Encoder(nn.Module):
                 layer_out = self._gradient_checkpointing_func(
                     layer_module.__call__,
                     hidden_states,
-                    None,
+                    position_mask,
                     T_patches,
                     H_patches,
                     W_patches,
@@ -686,7 +784,7 @@ class VJEPA21Encoder(nn.Module):
             else:
                 layer_out = layer_module(
                     hidden_states,
-                    position_mask=None,
+                    position_mask=position_mask,
                     T=T_patches,
                     H_patches=H_patches,
                     W_patches=W_patches,
@@ -834,13 +932,39 @@ class VJEPA21Predictor(nn.Module):
     def forward(
         self,
         encoder_hidden_states: torch.Tensor,
-        context_mask: list[torch.Tensor],
-        target_mask: list[torch.Tensor],
+        context_mask: Union[torch.Tensor, list[torch.Tensor]],
+        target_mask: Union[torch.Tensor, list[torch.Tensor]],
         mode: str = "video",
+        mask_index: int = 1,
+        context_is_masked: bool = False,
         output_attentions: bool = False,
         output_hidden_states: bool = False,
         **kwargs,
     ) -> VJEPA21PredictorOutput:
+        """
+        Args:
+            encoder_hidden_states: Encoder output. By default this is the *full*
+                token sequence and the predictor gathers the context itself at
+                `context_mask`. Set `context_is_masked=True` when passing an
+                encoder output that was already produced with
+                `encoder(..., masks=context_mask)`, which is the reference
+                convention.
+            mask_index: Which learnable mask token to inject. The reference uses
+                the index of the sequence-length group (`mask_index=i` in
+                `PredictorMultiSeqWrapper`); its default is 1.
+            mode: "video" or "img" ("image" is accepted as an alias of "img",
+                since the reference spells it that way).
+        """
+        if kwargs:
+            _warn_once(
+                f"VJEPA21Predictor.forward received unexpected keyword arguments "
+                f"{sorted(kwargs)}; they are ignored."
+            )
+
+        context_mask = _as_mask_list(context_mask)
+        target_mask = _as_mask_list(target_mask)
+        if context_mask is None or target_mask is None:
+            raise ValueError("the predictor requires both context_mask and target_mask")
         if len(context_mask) != len(target_mask):
             raise ValueError(
                 "context_mask and target_mask must have the same length, got "
@@ -852,10 +976,21 @@ class VJEPA21Predictor(nn.Module):
                 "pair. Call it once per mask pair instead."
             )
 
-        masked_states = apply_masks(encoder_hidden_states, context_mask)
+        if context_is_masked:
+            masked_states = encoder_hidden_states
+            if masked_states.shape[1] != context_mask[0].shape[1]:
+                raise ValueError(
+                    f"context_is_masked=True but the encoder output has "
+                    f"{masked_states.shape[1]} tokens while context_mask has "
+                    f"{context_mask[0].shape[1]}."
+                )
+        else:
+            masked_states = apply_masks(encoder_hidden_states, context_mask)
         _, N_ctxt, _ = masked_states.shape
 
-        hidden_states, position_masks = self.embeddings(masked_states, context_mask, target_mask)
+        hidden_states, position_masks = self.embeddings(
+            masked_states, context_mask, target_mask, mask_index=mask_index
+        )
 
         # Sort tokens by position so RoPE sees monotonically increasing ids
         argsort = torch.argsort(position_masks, dim=1)
@@ -864,13 +999,13 @@ class VJEPA21Predictor(nn.Module):
         position_masks = torch.gather(position_masks, 1, argsort.to(position_masks.device))
 
         if self.embeddings.img_mod_embed is not None:
-            if mode == "img":
+            if mode in ("img", "image"):
                 hidden_states = hidden_states + self.embeddings.img_mod_embed
             else:
                 hidden_states = hidden_states + self.embeddings.video_mod_embed
 
-        all_hidden_states = () if output_hidden_states else None
-        all_attentions = () if output_attentions else None
+        all_hidden_states: Optional[tuple] = () if output_hidden_states else None
+        all_attentions: Optional[tuple] = () if output_attentions else None
 
         for layer_module in self.layer:
             if output_hidden_states:
@@ -1036,6 +1171,14 @@ class VJEPA21PoolerCrossAttentionLayer(nn.Module):
 
 
 class VJEPA21AttentivePooler(nn.Module):
+    """Attentive pooler matching `AttentivePooler(depth=num_pooler_layers + 1)`.
+
+    The reference frozen-probe configs under `configs/eval_2_1/` use
+    `num_probe_blocks: 4` and `num_heads: 16`, i.e. three self-attention blocks
+    followed by one cross-attention block, with 16 heads. Those are the defaults
+    of `num_pooler_layers` and `num_pooler_heads`.
+    """
+
     def __init__(self, config: VJEPA21Config):
         super().__init__()
         self.query_tokens = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
@@ -1123,6 +1266,18 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
         out_layers=model.config.encoder_hierarchical_layers,
     ).multilevel_hidden_states
     ```
+
+    Example, the JEPA masked forward (encoder sees only the context tokens):
+
+    ```python
+    out = model(
+        pixel_values_videos,
+        masks=[context_idx],          # (B, K_ctx)
+        context_mask=[context_idx],
+        target_mask=[target_idx],     # (B, K_tgt)
+    )
+    prediction = out.predictor_output.last_hidden_state
+    ```
     """
 
     def __init__(self, config: VJEPA21Config):
@@ -1137,11 +1292,13 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
     def forward(
         self,
         pixel_values_videos: torch.Tensor,
-        context_mask: Optional[list[torch.Tensor]] = None,
-        target_mask: Optional[list[torch.Tensor]] = None,
+        masks: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
+        context_mask: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
+        target_mask: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
         skip_predictor: bool = False,
         return_hierarchical: bool = False,
         out_layers: Optional[list[int]] = None,
+        mask_index: int = 1,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         **kwargs,
@@ -1150,19 +1307,33 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
         Args:
             pixel_values_videos: Video tensor. Accepted layouts:
                 `(B, C, T, H, W)`, `(B, T, C, H, W)`, `(B, T, H, W, C)`.
+            masks: Optional list of `(B, K)` index tensors applied to the patch
+                tokens *before* the encoder layers. This reproduces the JEPA
+                training forward. When set, `context_mask` and `target_mask` must
+                be given explicitly, since the encoder output no longer spans the
+                full token grid.
             context_mask: List of `(B, K)` index tensors for context tokens.
-                Defaults to all tokens.
+                Defaults to all tokens when `masks` is None.
             target_mask: List of `(B, K)` index tensors for target tokens.
-                Defaults to all tokens.
+                Defaults to all tokens when `masks` is None.
             skip_predictor: Skip the predictor forward (encoder only).
             return_hierarchical: Return the concatenated distillation levels.
             out_layers: Encoder layer indices whose normalised features should be
                 returned in `multilevel_hidden_states`. Must be a subset of
                 `config.encoder_hierarchical_layers`.
+            mask_index: Which learnable predictor mask token to inject.
         """
         if pixel_values_videos is None:
             raise ValueError("pixel_values_videos is required")
 
+        unexpected = set(kwargs) - _SILENTLY_IGNORED_KWARGS
+        if unexpected:
+            _warn_once(
+                f"VJEPA21Model.forward received unexpected keyword arguments "
+                f"{sorted(unexpected)}; they are ignored. Accepted arguments: "
+                f"{sorted(_MODEL_FORWARD_KWARGS)}."
+            )
+
         output_attentions = (
             output_attentions if output_attentions is not None else self.config.output_attentions
         )
@@ -1172,6 +1343,10 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
             else self.config.output_hidden_states
         )
 
+        masks = _as_mask_list(masks)
+        context_mask = _as_mask_list(context_mask)
+        target_mask = _as_mask_list(target_mask)
+
         # When the predictor fuses several distillation levels (n_output_distillation > 1,
         # e.g. the ViT-g and ViT-G checkpoints) its input projection expects the
         # concatenated hierarchical features, not the last hidden state.
@@ -1181,11 +1356,11 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
 
         encoder_out = self.encoder(
             pixel_values_videos,
+            masks=masks,
             return_hierarchical=return_hierarchical or needs_hierarchical_input,
             out_layers=out_layers,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            **kwargs,
         )
         seq_output = encoder_out.last_hidden_state
 
@@ -1202,14 +1377,28 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
             num_tokens = seq_output.size(1)
             device = seq_output.device
 
-            if context_mask is None:
-                context_mask = [
-                    torch.arange(num_tokens, device=device).unsqueeze(0).expand(batch_size, -1)
-                ]
-            if target_mask is None:
-                target_mask = [
-                    torch.arange(num_tokens, device=device).unsqueeze(0).expand(batch_size, -1)
-                ]
+            if masks is not None:
+                # The encoder output only covers the context tokens, so the
+                # "all tokens" default is meaningless here.
+                if context_mask is None:
+                    context_mask = list(masks)
+                if target_mask is None:
+                    raise ValueError(
+                        "target_mask must be given explicitly when the encoder runs with "
+                        "`masks`, since the encoder output no longer spans the full token "
+                        "grid. Pass `skip_predictor=True` if you only want the encoder."
+                    )
+                context_is_masked = True
+            else:
+                context_is_masked = False
+                if context_mask is None:
+                    context_mask = [
+                        torch.arange(num_tokens, device=device).unsqueeze(0).expand(batch_size, -1)
+                    ]
+                if target_mask is None:
+                    target_mask = [
+                        torch.arange(num_tokens, device=device).unsqueeze(0).expand(batch_size, -1)
+                    ]
 
             mode = self._detect_mode(pixel_values_videos)
             predictor_output = self.predictor(
@@ -1217,10 +1406,16 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
                 context_mask,
                 target_mask,
                 mode=mode,
+                mask_index=mask_index,
+                context_is_masked=context_is_masked,
                 output_attentions=output_attentions,
                 output_hidden_states=output_hidden_states,
             )
-            masked_hidden_state = apply_masks(predictor_input, context_mask)
+            masked_hidden_state = (
+                predictor_input
+                if context_is_masked
+                else apply_masks(predictor_input, context_mask)
+            )
 
         return VJEPA21ModelOutput(
             last_hidden_state=seq_output,
@@ -1247,7 +1442,11 @@ class VJEPA21Model(VJEPA21PreTrainedModel):
 
 
 class VJEPA21ForVideoClassification(VJEPA21PreTrainedModel):
-    """V-JEPA 2.1 with attentive pooler + classification head."""
+    """V-JEPA 2.1 with attentive pooler + classification head.
+
+    The pooler and the classifier are always randomly initialised: this class is
+    the frozen-probe / fine-tuning entry point, not a pretrained classifier.
+    """
 
     def __init__(self, config: VJEPA21Config):
         super().__init__(config)
@@ -1266,16 +1465,30 @@ class VJEPA21ForVideoClassification(VJEPA21PreTrainedModel):
         **kwargs,
     ) -> ImageClassifierOutput:
         r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the classification loss. Indices should be in
-            `[0, ..., config.num_labels - 1]`.
+        labels (`torch.LongTensor` of shape `(batch_size,)` or `(batch_size, num_labels)`, *optional*):
+            Labels for computing the classification loss. Integer indices in
+            `[0, ..., config.num_labels - 1]` give single-label classification;
+            a float multi-hot tensor gives multi-label classification (BCE), which
+            is what a dataset such as XD-Violence needs. Set
+            `config.problem_type` explicitly to remove the ambiguity.
         """
+        # Only forward the arguments the backbone understands; the Trainer injects
+        # extras such as `num_items_in_batch` that would otherwise reach the encoder.
+        forwarded = {k: v for k, v in kwargs.items() if k in _MODEL_FORWARD_KWARGS}
+        ignored = set(kwargs) - set(forwarded) - _SILENTLY_IGNORED_KWARGS
+        if ignored:
+            _warn_once(
+                f"VJEPA21ForVideoClassification.forward received unexpected keyword "
+                f"arguments {sorted(ignored)}; they are ignored."
+            )
+        forwarded.pop("skip_predictor", None)
+
         outputs = self.vjepa21(
             pixel_values_videos,
             skip_predictor=True,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            **kwargs,
+            **forwarded,
         )
         pooled = self.pooler(outputs.last_hidden_state)
         logits = self.classifier(pooled)

tests/test_end_to_end_pipeline.py

@@ -30,25 +30,22 @@ PORT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 VJEPA2_REPO = os.environ.get("VJEPA2_REPO", "")
 CKPT = os.environ.get("VJEPA21_CKPT", "")
 
-pytestmark = pytest.mark.skipif(
-    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "src", "datasets")),
-    reason="set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
-)
-
 
-def _processor():
+def _processor(**kwargs):
     name = "_vjepa21_port_under_test"
     if name not in sys.modules:
         pkg = types.ModuleType(name)
         pkg.__path__ = [PORT_DIR]
         sys.modules[name] = pkg
     mod = importlib.import_module(f"{name}.video_processing_vjepa21")
-    return mod.VJEPA21VideoProcessor()
+    return mod.VJEPA21VideoProcessor(**kwargs)
 
 
 def _reference_transform(clip, crop=384):
     """Reference `EvalVideoTransform`, centre crop instead of the sliding views."""
-    sys.path.insert(0, VJEPA2_REPO)
+    repo = os.path.abspath(VJEPA2_REPO)
+    if repo not in sys.path:
+        sys.path.insert(0, repo)
     import src.datasets.utils.video.transforms as video_transforms
     import src.datasets.utils.video.volume_transforms as volume_transforms
 
@@ -72,7 +69,61 @@ def _clip(height, width, frames=4, seed=0):
     return [rng.integers(0, 255, (height, width, 3), dtype=np.uint8) for _ in range(frames)]
 
 
-# --- pixel level ------------------------------------------------------------
+# --- 1. antialiasing: no reference repository needed -------------------------
+
+
+def test_resize_defaults_to_no_antialiasing():
+    """Regression guard for the resize backend.
+
+    This used to be run on `torch.zeros(1, 3, 1080, 1920)`, which cannot fail:
+    resizing a constant image gives the same result with and without
+    antialiasing, so `torch.equal(default, forced)` held regardless of the
+    default. It needs real content, and it needs to show that the antialiased
+    path is genuinely different — otherwise the guard is still vacuous.
+    """
+    processor = _processor()
+    torch.manual_seed(0)
+    big = torch.rand(1, 3, 1080, 1920)
+    size = type(processor.size)(shortest_edge=384)
+
+    default = processor.resize(big, size)
+    forced_off = processor.resize(big, size, antialias=False)
+    forced_on = processor.resize(big, size, antialias=True)
+
+    assert torch.equal(default, forced_off), "resize() must default to antialias=False"
+    gap = (forced_on - default).abs().max().item()
+    print(f"\n[antialias] max|Δ| antialiased vs not = {gap:.4f} "
+          f"({gap * 255:.1f}/255 on an 8-bit scale)")
+    assert gap > 1e-3, (
+        "antialiased and non-antialiased downscaling are indistinguishable; "
+        "the guard is not testing anything"
+    )
+
+
+def test_antialiasing_gap_is_negligible_when_upscaling():
+    """Explains why this only ever mattered for downscaling: on a 240p source
+    resized up to 384 the two backends agree, which is how the bug survived."""
+    processor = _processor()
+    torch.manual_seed(0)
+    small = torch.rand(1, 3, 240, 320)
+    size = type(processor.size)(shortest_edge=384)
+    gap = (
+        processor.resize(small, size, antialias=True)
+        - processor.resize(small, size, antialias=False)
+    ).abs().max().item()
+    print(f"[antialias, upscale] max|Δ| = {gap:.6f}")
+    assert gap < 1e-5
+
+
+def test_processor_geometry_is_independent_of_aspect_ratio():
+    """Short side to exactly `crop_size`, then a square centre crop."""
+    processor = _processor()
+    for h, w in ((240, 320), (320, 240), (1080, 1920), (400, 400)):
+        pv = processor([_clip(h, w, frames=2)], return_tensors="pt")["pixel_values_videos"]
+        assert pv.shape == (1, 2, 3, 384, 384), (h, w, pv.shape)
+
+
+# --- 2. pixel level against the reference transform --------------------------
 
 RESOLUTIONS = [
     pytest.param(240, 320, id="240p-upscale"),
@@ -81,7 +132,13 @@ RESOLUTIONS = [
     pytest.param(1080, 1920, id="1080p-downscale"),
 ]
 
+needs_reference = pytest.mark.skipif(
+    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "src", "datasets")),
+    reason="set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
+)
+
 
+@needs_reference
 @pytest.mark.parametrize("height,width", RESOLUTIONS)
 def test_processor_matches_reference_transform(height, width):
     """The processor must agree with the reference transform to within the 8-bit
@@ -108,18 +165,36 @@ def test_processor_matches_reference_transform(height, width):
     assert diff.max().item() <= quantisation_floor * 1.05
 
 
-def test_antialiasing_is_disabled():
-    """Regression guard for the resize backend, independent of resolution."""
+@needs_reference
+def test_antialiased_processor_would_fail_the_reference_comparison():
+    """Shows the tolerance above is tight enough to catch the regression it was
+    written for, rather than passing either way."""
+    clip = _clip(1080, 1920)
+    reference = _reference_transform(clip)
+
     processor = _processor()
-    big = torch.zeros(1, 3, 1080, 1920)
-    default = processor.resize(big, type(processor.size)(shortest_edge=384))
-    forced = processor.resize(big, type(processor.size)(shortest_edge=384), antialias=False)
-    assert torch.equal(default, forced), "resize() must default to antialias=False"
+    original_resize = type(processor).resize
+
+    def antialiased(self, image, size, **kwargs):
+        kwargs["antialias"] = True
+        return super(type(self), self).resize(image, size, **kwargs)
+
+    try:
+        type(processor).resize = antialiased
+        got = processor([clip], return_tensors="pt")["pixel_values_videos"][0]
+        got = got.permute(1, 0, 2, 3)
+    finally:
+        type(processor).resize = original_resize
+
+    gap = (reference - got).abs().max().item()
+    print(f"\n[antialiased 1080p] max|Δ| vs reference = {gap:.4e}")
+    assert gap > 1.0 / 255.0 / 0.229 * 1.05
 
 
-# --- feature level ----------------------------------------------------------
+# --- 3. feature level --------------------------------------------------------
 
 
+@needs_reference
 @pytest.mark.skipif(not CKPT, reason="set VJEPA21_CKPT to run the model-level comparison")
 @torch.no_grad()
 @pytest.mark.parametrize("height,width", RESOLUTIONS)

tests/test_parity_official.py

@@ -1,15 +1,19 @@
-"""Numerical parity test: HF port vs. the reference V-JEPA 2.1 implementation.
-
-Builds the reference encoder from a local clone of ``facebookresearch/vjepa2``
-and the ported ``VJEPA21Model``, copies the reference weights into the port and
-compares the outputs.
-
-Usage:
-    git clone https://github.com/facebookresearch/vjepa2.git
-    pip install timm einops
-    VJEPA2_REPO=/path/to/vjepa2 python -m pytest test_parity_official.py -s
-
-The test is skipped when the reference repository is not available.
+"""Implementation parity against the reference `facebookresearch/vjepa2`.
+
+Reference modules are built directly from `app/vjepa_2_1/models/`, their weights
+are copied into the port, and the outputs are compared. Reduced hidden sizes keep
+the suite fast: these checks establish that the *code paths* are equivalent,
+which is width-independent. Parity on the published weights at full resolution is
+a separate concern, covered by `verify_vjepa21_port.py`.
+
+    VJEPA2_REPO=/path/to/vjepa2 python -m pytest test_parity_official.py -s -q
+
+Note on tolerances. The encoder is compared eager-vs-eager and should be
+bit-exact. The predictor is not: `VisionTransformerPredictor.__init__` has no
+`use_sdpa` parameter, so it is swallowed by `**kwargs` and the reference
+predictor blocks always run SDPA, while the port runs eager here. That plus the
+gather-vs-stack reordering puts the residual around 1e-06, three orders of
+magnitude inside the 1e-3 tolerance Meta uses for its own ports.
 """
 
 from __future__ import annotations
@@ -18,314 +22,390 @@ import importlib
 import os
 import sys
 import types
+from functools import partial
 
 import pytest
 import torch
+import torch.nn as nn
 
-VJEPA2_REPO = os.environ.get("VJEPA2_REPO", "")
 PORT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+VJEPA2_REPO = os.environ.get("VJEPA2_REPO", "")
+
+pytestmark = pytest.mark.skipif(
+    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "app", "vjepa_2_1")),
+    reason="set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
+)
 
 
 def _load_port():
-    """Import the port's modules as a package so relative imports resolve."""
     name = "_vjepa21_port_under_test"
     if name not in sys.modules:
         pkg = types.ModuleType(name)
         pkg.__path__ = [PORT_DIR]
         sys.modules[name] = pkg
-    cfg_mod = importlib.import_module(f"{name}.configuration_vjepa21")
-    mdl_mod = importlib.import_module(f"{name}.modeling_vjepa21")
-    return cfg_mod.VJEPA21Config, mdl_mod.VJEPA21Model
+    cfg = importlib.import_module(f"{name}.configuration_vjepa21")
+    mdl = importlib.import_module(f"{name}.modeling_vjepa21")
+    return cfg.VJEPA21Config, mdl.VJEPA21Model, mdl.apply_masks
+
+
+VJEPA21Config, VJEPA21Model, apply_masks = _load_port()
+
+
+def _reference_classes():
+    repo = os.path.abspath(VJEPA2_REPO)
+    if repo not in sys.path:
+        sys.path.insert(0, repo)
+    from app.vjepa_2_1.models.predictor import VisionTransformerPredictor
+    from app.vjepa_2_1.models.vision_transformer import VisionTransformer
+
+    return VisionTransformer, VisionTransformerPredictor
 
-pytestmark = pytest.mark.skipif(
-    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "app", "vjepa_2_1")),
-    reason="Reference repository not found; set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
-)
 
-# --- reference -> port state dict mapping -----------------------------------
+# Small but structurally faithful: 64x64 x 4 frames, patch 16, tubelet 2
+# -> 2 x 4 x 4 = 32 tokens. head_dim = 96 // 6 = 16, so the RoPE split is
+# d=h=w=4 with 4 unrotated dimensions, exercising the same branch as the real
+# checkpoints (head_dim 64, split 20/20/20 + 4).
+DIM, HEADS, DEPTH = 96, 6, 12
+PRED_DIM, PRED_HEADS, PRED_DEPTH = 48, 6, 12
+IMG, PATCH, FRAMES, TUBELET = 64, 16, 4, 2
+NORM = partial(nn.LayerNorm, eps=1e-6)
 
 
-def convert_reference_encoder_state_dict(sd: dict, hidden_size: int) -> dict:
-    """Map reference encoder parameter names onto the ported encoder."""
+# ---------------------------------------------------------------------------
+# weight mapping
+# ---------------------------------------------------------------------------
+
+
+def _map_block(prefix, idx, sub, tensor, hidden, out):
+    if sub.startswith("attn.qkv."):
+        kind = sub.rsplit(".", 1)[-1]
+        q, k, v = tensor.split(hidden, dim=0)
+        out[f"{prefix}.layer.{idx}.attention.query.{kind}"] = q
+        out[f"{prefix}.layer.{idx}.attention.key.{kind}"] = k
+        out[f"{prefix}.layer.{idx}.attention.value.{kind}"] = v
+    elif sub.startswith("attn.proj."):
+        out[f"{prefix}.layer.{idx}.attention.proj." + sub.rsplit(".", 1)[-1]] = tensor
+    else:
+        out[f"{prefix}.layer.{idx}.{sub}"] = tensor
+
+
+def reference_to_port(enc_sd, pred_sd, hidden, pred_hidden):
+    """Rename reference tensors to port names. Splits the fused QKV projection."""
     out = {}
-    for k, v in sd.items():
-        if k == "img_mod_embed":
-            out["embeddings.img_mod_embed"] = v
-        elif k == "video_mod_embed":
-            out["embeddings.video_mod_embed"] = v
+    for k, v in enc_sd.items():
+        if k in ("img_mod_embed", "video_mod_embed"):
+            out[f"encoder.embeddings.{k}"] = v
         elif k.startswith("patch_embed_img."):
-            out["embeddings.patch_embeddings_img." + k[len("patch_embed_img.") :]] = v
+            out["encoder.embeddings.patch_embeddings_img." + k[len("patch_embed_img."):]] = v
         elif k.startswith("patch_embed."):
-            out["embeddings.patch_embeddings." + k[len("patch_embed.") :]] = v
+            out["encoder.embeddings.patch_embeddings." + k[len("patch_embed."):]] = v
         elif k.startswith("norms_block."):
-            out[k] = v
+            out["encoder." + k] = v
         elif k.startswith("blocks."):
-            rest = k[len("blocks.") :]
-            idx, sub = rest.split(".", 1)
-            if sub.startswith("attn.qkv."):
-                kind = sub.split(".")[-1]  # weight | bias
-                # Reference packs qkv as [q; k; v] along dim 0.
-                q, key, val = v.split(hidden_size, dim=0)
-                out[f"layer.{idx}.attention.query.{kind}"] = q
-                out[f"layer.{idx}.attention.key.{kind}"] = key
-                out[f"layer.{idx}.attention.value.{kind}"] = val
-            elif sub.startswith("attn.proj."):
-                out[f"layer.{idx}.attention.proj." + sub.split(".")[-1]] = v
-            else:
-                out[f"layer.{idx}.{sub}"] = v
+            idx, sub = k[len("blocks."):].split(".", 1)
+            _map_block("encoder", idx, sub, v, hidden, out)
+        elif k == "pos_embed":
+            continue  # unused under RoPE
         else:
-            raise KeyError(f"Unmapped reference key: {k}")
+            raise AssertionError(f"unmapped reference encoder key: {k}")
+
+    for k, v in pred_sd.items():
+        if k in ("img_mod_embed", "video_mod_embed"):
+            out[f"predictor.embeddings.{k}"] = v
+        elif k.startswith("predictor_embed."):
+            out["predictor.embeddings.predictor_embed." + k[len("predictor_embed."):]] = v
+        elif k.startswith("mask_tokens."):
+            out["predictor.embeddings." + k] = v
+        elif k.startswith("predictor_norm."):
+            out["predictor.layernorm." + k[len("predictor_norm."):]] = v
+        elif k.startswith("predictor_proj_context."):
+            out["predictor.proj_context." + k[len("predictor_proj_context."):]] = v
+        elif k.startswith("predictor_proj."):
+            out["predictor.proj." + k[len("predictor_proj."):]] = v
+        elif k.startswith("predictor_blocks."):
+            idx, sub = k[len("predictor_blocks."):].split(".", 1)
+            _map_block("predictor", idx, sub, v, pred_hidden, out)
+        elif k == "predictor_pos_embed":
+            continue
+        else:
+            raise AssertionError(f"unmapped reference predictor key: {k}")
     return out
 
 
-def _build_pair(n_output_distillation: int = 1, **overrides):
-    sys.path.insert(0, VJEPA2_REPO)
-    from app.vjepa_2_1.models.vision_transformer import vit_base
-
-    VJEPA21Config, VJEPA21Model = _load_port()
-
-    kw = dict(
-        img_size=(64, 64),
-        patch_size=16,
-        num_frames=4,
-        tubelet_size=2,
-        uniform_power=True,
-        use_rope=True,
-        img_temporal_dim_size=1,
-        interpolate_rope=True,
-        modality_embedding=True,
-        n_output_distillation=n_output_distillation,
-        use_sdpa=False,
-        use_silu=False,
-        wide_silu=True,
-    )
-    kw.update(overrides)
+# ---------------------------------------------------------------------------
+# builders
+# ---------------------------------------------------------------------------
+
 
+def build_reference(n_distill=1, out_layers=None, teacher_embed_dim=None):
+    VisionTransformer, VisionTransformerPredictor = _reference_classes()
     torch.manual_seed(0)
-    ref = vit_base(**kw).eval()
+    encoder = VisionTransformer(
+        img_size=(IMG, IMG), patch_size=PATCH, num_frames=FRAMES, tubelet_size=TUBELET,
+        in_chans=3, embed_dim=DIM, depth=DEPTH, num_heads=HEADS, mlp_ratio=4.0,
+        qkv_bias=True, norm_layer=NORM, use_rope=True, use_sdpa=False,
+        uniform_power=False, handle_nonsquare_inputs=True, img_temporal_dim_size=1,
+        interpolate_rope=True, modality_embedding=True,
+        n_output_distillation=n_distill, out_layers=out_layers,
+    ).eval()
+    predictor = VisionTransformerPredictor(
+        img_size=(IMG, IMG), patch_size=PATCH, num_frames=FRAMES, tubelet_size=TUBELET,
+        embed_dim=DIM, predictor_embed_dim=PRED_DIM, depth=PRED_DEPTH,
+        num_heads=PRED_HEADS, mlp_ratio=4.0, qkv_bias=True, norm_layer=NORM,
+        use_mask_tokens=True, num_mask_tokens=8, zero_init_mask_tokens=True,
+        use_silu=False, wide_silu=True, use_rope=True, interpolate_rope=True,
+        modality_embedding=True, img_temporal_dim_size=1, uniform_power=False,
+        return_all_tokens=True, teacher_embed_dim=teacher_embed_dim,
+        n_output_distillation=n_distill,
+    ).eval()
+    return encoder, predictor
+
 
+def build_port(n_distill=1, teacher_embed_dim=None):
     cfg = VJEPA21Config(
-        patch_size=kw["patch_size"],
-        crop_size=kw["img_size"][0],
-        tubelet_size=kw["tubelet_size"],
-        hidden_size=768,
-        num_attention_heads=12,
-        num_hidden_layers=12,
-        hidden_act="silu" if kw["use_silu"] else "gelu",
-        wide_silu=kw["wide_silu"],
-        img_temporal_dim_size=kw["img_temporal_dim_size"],
-        interpolate_rope=kw["interpolate_rope"],
-        modality_embedding=kw["modality_embedding"],
-        n_output_distillation=kw["n_output_distillation"],
-        pred_hidden_size=384,
-        pred_num_attention_heads=12,
-        pred_num_hidden_layers=12,
-        pred_teacher_embed_dim=1664 if n_output_distillation == 1 else None,
-        pred_return_all_tokens=True,
+        patch_size=PATCH, crop_size=IMG, frames_per_clip=FRAMES, tubelet_size=TUBELET,
+        hidden_size=DIM, num_attention_heads=HEADS, num_hidden_layers=DEPTH,
+        mlp_ratio=4.0, num_pooler_heads=HEADS,
+        pred_hidden_size=PRED_DIM, pred_num_attention_heads=PRED_HEADS,
+        pred_num_hidden_layers=PRED_DEPTH, pred_num_mask_tokens=8,
+        pred_teacher_embed_dim=teacher_embed_dim, pred_return_all_tokens=True,
+        n_output_distillation=n_distill, img_temporal_dim_size=1,
+        interpolate_rope=True, modality_embedding=True,
         attn_implementation="eager",
     )
-    port = VJEPA21Model(cfg).eval()
+    return VJEPA21Model(cfg).eval()
+
+
+def make_pair(n_distill=1, out_layers=None, teacher_embed_dim=None):
+    """Reference pair plus a port loaded with exactly those weights."""
+    ref_enc, ref_pred = build_reference(n_distill, out_layers, teacher_embed_dim)
+    port = build_port(n_distill, teacher_embed_dim)
+    state = reference_to_port(ref_enc.state_dict(), ref_pred.state_dict(), DIM, PRED_DIM)
+    missing, unexpected = port.load_state_dict(state, strict=False)
+    # This is the check that matters: a parameter with no counterpart stays at its
+    # random initialisation and nothing else in the suite would notice.
+    assert not missing, f"port parameters with no reference origin: {sorted(missing)[:10]}"
+    assert not unexpected, f"reference tensors the port has no home for: {sorted(unexpected)[:10]}"
+    return ref_enc, ref_pred, port
+
+
+def maxdiff(a, b):
+    return (a - b).abs().max().item()
+
+
+@pytest.fixture(scope="module")
+def video():
+    torch.manual_seed(1)
+    return torch.randn(2, 3, FRAMES, IMG, IMG)
 
-    converted = convert_reference_encoder_state_dict(ref.state_dict(), cfg.hidden_size)
-    missing, unexpected = port.encoder.load_state_dict(converted, strict=False)
-    assert not unexpected, f"unexpected keys: {unexpected}"
-    assert not missing, f"missing keys: {missing}"
-    return ref, port
+
+# ---------------------------------------------------------------------------
+# encoder
+# ---------------------------------------------------------------------------
 
 
 @torch.no_grad()
 @pytest.mark.parametrize("n_distill", [1, 4])
-def test_last_hidden_state_matches_reference(n_distill):
-    """n_distill=1 is the ViT-B/ViT-L regime, n_distill=4 the ViT-g/ViT-G one."""
-    ref, port = _build_pair(n_distill)
-    x = torch.randn(2, 3, 4, 64, 64)
+def test_encoder_video_branch(video, n_distill):
+    ref_enc, _, port = make_pair(n_distill=n_distill)
+    expected = ref_enc(video)
+    got = port(pixel_values_videos=video, skip_predictor=True).last_hidden_state
+    d = maxdiff(expected, got)
+    print(f"\n[encoder n_distill={n_distill}] max|Δ| = {d:.3e}")
+    assert d < 1e-5
 
-    expected = ref(x)
-    got = port(pixel_values_videos=x, skip_predictor=True).last_hidden_state
 
-    assert expected.shape == got.shape
-    max_abs = (expected - got).abs().max().item()
-    print(f"\n[video n_distill={n_distill}] max |diff| = {max_abs:.3e}")
-    assert max_abs < 1e-5
+@torch.no_grad()
+def test_encoder_hierarchical_output(video):
+    ref_enc, _, port = make_pair(n_distill=4)
+    expected = ref_enc(video, training=True)          # concatenated levels
+    got = port(
+        pixel_values_videos=video, skip_predictor=True, return_hierarchical=True
+    ).hierarchical_hidden_state
+    assert expected.shape == got.shape == (2, 32, DIM * 4)
+    d = maxdiff(expected, got)
+    print(f"[encoder hierarchical] max|Δ| = {d:.3e}")
+    assert d < 1e-5
 
 
 @torch.no_grad()
-def test_image_branch_matches_reference():
-    """T == img_temporal_dim_size routes through the image patch embedding."""
-    ref, port = _build_pair()
-    x = torch.randn(2, 3, 1, 64, 64)
+def test_encoder_image_branch():
+    ref_enc, _, port = make_pair()
+    torch.manual_seed(2)
+    image = torch.randn(2, 3, 1, IMG, IMG)
+    expected = ref_enc(image)
+    got = port(pixel_values_videos=image, skip_predictor=True).last_hidden_state
+    assert got.shape == (2, 16, DIM)
+    d = maxdiff(expected, got)
+    print(f"[encoder image branch] max|Δ| = {d:.3e}")
+    assert d < 1e-5
 
-    expected = ref(x)
-    got = port(pixel_values_videos=x, skip_predictor=True).last_hidden_state
 
-    max_abs = (expected - got).abs().max().item()
-    print(f"[image] max |diff| = {max_abs:.3e}")
-    assert max_abs < 1e-5
+@torch.no_grad()
+def test_encoder_non_square_input():
+    """Exercises the RoPE interpolation on both spatial axes independently."""
+    ref_enc, _, port = make_pair()
+    torch.manual_seed(3)
+    x = torch.randn(1, 3, FRAMES, 96, 128)
+    expected = ref_enc(x)
+    got = port(pixel_values_videos=x, skip_predictor=True).last_hidden_state
+    assert got.shape == (1, 2 * 6 * 8, DIM)
+    d = maxdiff(expected, got)
+    print(f"[encoder 96x128] max|Δ| = {d:.3e}")
+    assert d < 1e-5
 
 
 @torch.no_grad()
 @pytest.mark.parametrize("n_distill", [1, 4])
-def test_multilevel_matches_reference_out_layers(n_distill):
-    """`out_layers` must reproduce the reference multi-level probe features."""
-    ref, port = _build_pair(n_distill)
-    ref.out_layers = [2, 5, 8, 11]
-    x = torch.randn(2, 3, 4, 64, 64)
-
-    expected = ref(x)  # list of per-level normalised features
+def test_encoder_out_layers(video, n_distill):
+    """`out_layers` is a constructor argument on the reference and a forward
+    argument on the port; the per-level norms must still agree."""
+    layers = [2, 5, 8, 11]
+    ref_enc, _, port = make_pair(n_distill=n_distill, out_layers=layers)
+    expected = ref_enc(video)                          # list, one tensor per level
     got = port(
-        pixel_values_videos=x, skip_predictor=True, out_layers=[2, 5, 8, 11]
+        pixel_values_videos=video, skip_predictor=True, out_layers=layers
     ).multilevel_hidden_states
-
     assert len(expected) == len(got) == 4
-    for i, (e, g) in enumerate(zip(expected, got)):
-        max_abs = (e - g).abs().max().item()
-        print(f"[level {[2, 5, 8, 11][i]}] max |diff| = {max_abs:.3e}")
-        assert max_abs < 1e-5
+    for i, (a, b) in enumerate(zip(expected, got)):
+        d = maxdiff(a, b)
+        print(f"[out_layers n_distill={n_distill}] level {layers[i]}: max|Δ| = {d:.3e}")
+        assert d < 1e-5
 
 
 @torch.no_grad()
-def test_non_square_resolution_matches_reference():
-    """RoPE interpolation must agree on a non-native, non-square input."""
-    ref, port = _build_pair()
-    x = torch.randn(1, 3, 4, 96, 128)
+def test_masked_encoder_forward(video):
+    """The JEPA training forward: tokens are dropped before the blocks and RoPE
+    receives their original indices. Without this the predictor can only ever be
+    fed representations computed with attention over tokens that were masked
+    during pre-training.
+    """
+    ref_enc, _, port = make_pair()
+    torch.manual_seed(4)
+    idx = torch.stack([torch.randperm(32)[:14] for _ in range(2)])
+    expected = ref_enc(video, masks=[idx])
+    got = port.encoder(video, masks=[idx]).last_hidden_state
+    assert got.shape == (2, 14, DIM)
+    d = maxdiff(expected, got)
+    print(f"[masked encoder] max|Δ| = {d:.3e}")
+    assert d < 1e-5
 
-    expected = ref(x)
-    got = port(pixel_values_videos=x, skip_predictor=True).last_hidden_state
 
-    max_abs = (expected - got).abs().max().item()
-    print(f"[96x128] max |diff| = {max_abs:.3e}")
-    assert max_abs < 1e-5
+@torch.no_grad()
+def test_masked_encoder_is_not_equivalent_to_masking_the_output(video):
+    """Control: gathering after the blocks gives a different answer, so the test
+    above is checking something real."""
+    ref_enc, _, port = make_pair()
+    idx = torch.arange(0, 14).unsqueeze(0).expand(2, -1)
+    inside = port.encoder(video, masks=[idx]).last_hidden_state
+    outside = apply_masks(
+        port(pixel_values_videos=video, skip_predictor=True).last_hidden_state, [idx]
+    )
+    assert not torch.allclose(inside, outside, atol=1e-4)
 
 
-# --- predictor ---------------------------------------------------------------
+# ---------------------------------------------------------------------------
+# predictor
+# ---------------------------------------------------------------------------
 
 
-def convert_reference_predictor_state_dict(sd: dict, pred_hidden_size: int) -> dict:
-    """Map reference predictor parameter names onto the ported predictor."""
-    out = {}
-    for k, v in sd.items():
-        if k in ("img_mod_embed", "video_mod_embed"):
-            out[f"embeddings.{k}"] = v
-        elif k.startswith("predictor_embed."):
-            out["embeddings.predictor_embed." + k[len("predictor_embed.") :]] = v
-        elif k.startswith("mask_tokens."):
-            out["embeddings." + k] = v
-        elif k.startswith("predictor_norm."):
-            out["layernorm." + k[len("predictor_norm.") :]] = v
-        elif k.startswith("predictor_proj_context."):
-            out["proj_context." + k[len("predictor_proj_context.") :]] = v
-        elif k.startswith("predictor_proj."):
-            out["proj." + k[len("predictor_proj.") :]] = v
-        elif k.startswith("predictor_blocks."):
-            rest = k[len("predictor_blocks.") :]
-            idx, sub = rest.split(".", 1)
-            if sub.startswith("attn.qkv."):
-                kind = sub.split(".")[-1]
-                q, key, val = v.split(pred_hidden_size, dim=0)
-                out[f"layer.{idx}.attention.query.{kind}"] = q
-                out[f"layer.{idx}.attention.key.{kind}"] = key
-                out[f"layer.{idx}.attention.value.{kind}"] = val
-            elif sub.startswith("attn.proj."):
-                out[f"layer.{idx}.attention.proj." + sub.split(".")[-1]] = v
-            else:
-                out[f"layer.{idx}.{sub}"] = v
-        else:
-            raise KeyError(f"Unmapped reference predictor key: {k}")
-    return out
+def _masks(batch=2, n_tokens=32, n_ctx=20):
+    ctx = torch.arange(0, n_ctx).unsqueeze(0).expand(batch, -1)
+    tgt = torch.arange(n_ctx, n_tokens).unsqueeze(0).expand(batch, -1)
+    return ctx, tgt
 
 
 @torch.no_grad()
-def test_predictor_matches_reference():
-    sys.path.insert(0, VJEPA2_REPO)
-    from app.vjepa_2_1.models.predictor import vit_predictor
-
-    _, port = _build_pair()
-
-    torch.manual_seed(1)
-    ref = vit_predictor(
-        img_size=(64, 64), patch_size=16, use_mask_tokens=True, embed_dim=768,
-        predictor_embed_dim=384, teacher_embed_dim=1664, num_frames=4, tubelet_size=2,
-        depth=12, num_heads=12, num_mask_tokens=8, use_rope=True, uniform_power=False,
-        use_sdpa=False, use_silu=False, wide_silu=True, n_output_distillation=1,
-        return_all_tokens=True, img_temporal_dim_size=1, modality_embedding=True,
-        zero_init_mask_tokens=True, interpolate_rope=True,
-    ).eval()
-
-    converted = convert_reference_predictor_state_dict(ref.state_dict(), 384)
-    missing, unexpected = port.predictor.load_state_dict(converted, strict=False)
-    assert not unexpected, f"unexpected keys: {unexpected}"
-    assert not missing, f"missing keys: {missing}"
-
-    B, N, D = 2, 32, 768
-    z = torch.randn(B, N, D)
-    ctx = torch.arange(0, 20).unsqueeze(0).expand(B, -1)
-    tgt = torch.arange(20, N).unsqueeze(0).expand(B, -1)
+@pytest.mark.parametrize(
+    "n_distill,teacher", [(1, DIM), (1, None), (4, None), (4, DIM * 4)]
+)
+def test_predictor(video, n_distill, teacher):
+    ref_enc, ref_pred, port = make_pair(n_distill=n_distill, teacher_embed_dim=teacher)
+    ctx, tgt = _masks()
 
-    from _vjepa21_port_under_test.modeling_vjepa21 import apply_masks
+    # The predictor input is the concatenated levels when the encoder distils
+    # more than one, and the last hidden state otherwise.
+    z = ref_enc(video, training=True) if n_distill > 1 else ref_enc(video)
 
-    ref_pred, ref_ctx = ref(apply_masks(z, [ctx]), [ctx], [tgt], mod="video")
+    exp_pred, exp_ctx = ref_pred(apply_masks(z, [ctx]), [ctx], [tgt], mod="video")
     got = port.predictor(z, [ctx], [tgt], mode="video")
 
-    d_pred = (ref_pred - got.last_hidden_state).abs().max().item()
-    d_ctx = (ref_ctx - got.context_hidden_state).abs().max().item()
-    print(f"[predictor target] max |diff| = {d_pred:.3e}")
-    print(f"[predictor context] max |diff| = {d_ctx:.3e}")
-    assert d_pred < 1e-5 and d_ctx < 1e-5
+    dp, dc = maxdiff(exp_pred, got.last_hidden_state), maxdiff(exp_ctx, got.context_hidden_state)
+    print(f"\n[predictor n_distill={n_distill} teacher={teacher}] "
+          f"target max|Δ| = {dp:.3e}   context max|Δ| = {dc:.3e}")
+    assert dp < 1e-3 and dc < 1e-3
 
 
 @torch.no_grad()
-def test_hierarchical_output_matches_reference():
-    """With n_output_distillation=4 the encoder returns the concatenated levels.
-    This is what the ViT-g and ViT-G predictors consume."""
-    ref, port = _build_pair(4)
-    x = torch.randn(2, 3, 4, 64, 64)
-
-    expected = ref(x, training=True)  # reference returns cat(hier) in training mode
-    got = port(
-        pixel_values_videos=x, skip_predictor=True, return_hierarchical=True
-    ).hierarchical_hidden_state
-
-    assert expected.shape == got.shape == (2, 32, 768 * 4)
-    max_abs = (expected - got).abs().max().item()
-    print(f"[hierarchical output] max |diff| = {max_abs:.3e}")
-    assert max_abs < 1e-5
+def test_predictor_with_trained_mask_tokens(video):
+    """The shipped checkpoints have `zero_init_mask_tokens=True`, so every mask
+    token is exactly zero and any comparison of the mask-token path passes even
+    if the lookup is broken. Give them distinct values first, then compare.
+    """
+    ref_enc, ref_pred, port = make_pair()
+    torch.manual_seed(5)
+    for i, token in enumerate(ref_pred.mask_tokens):
+        nn.init.normal_(token, std=0.1 * (i + 1))
+    # copy the same values into the port
+    port_tokens = port.predictor.embeddings.mask_tokens
+    for dst, src in zip(port_tokens, ref_pred.mask_tokens):
+        dst.copy_(src)
+    assert max(t.abs().max().item() for t in ref_pred.mask_tokens) > 0
+
+    ctx, tgt = _masks()
+    z = ref_enc(video)
+    for mask_index in (0, 1, 3, 7, 9):  # 9 exercises the modulo wrap
+        exp_pred, exp_ctx = ref_pred(
+            apply_masks(z, [ctx]), [ctx], [tgt], mod="video", mask_index=mask_index
+        )
+        got = port.predictor(z, [ctx], [tgt], mode="video", mask_index=mask_index)
+        dp = maxdiff(exp_pred, got.last_hidden_state)
+        dc = maxdiff(exp_ctx, got.context_hidden_state)
+        print(f"[mask_index={mask_index}] target max|Δ| = {dp:.3e}  context max|Δ| = {dc:.3e}")
+        assert dp < 1e-3 and dc < 1e-3
+
+    # and the indices must actually differ from one another
+    a = port.predictor(z, [ctx], [tgt], mask_index=0).last_hidden_state
+    b = port.predictor(z, [ctx], [tgt], mask_index=1).last_hidden_state
+    assert not torch.allclose(a, b, atol=1e-4)
 
 
 @torch.no_grad()
-def test_hierarchical_predictor_matches_reference():
-    """ViT-g / ViT-G predictor: 2-layer MLP over the concatenated levels."""
-    sys.path.insert(0, VJEPA2_REPO)
-    from app.vjepa_2_1.models.predictor import vit_predictor
-
-    _, port = _build_pair(4)
-
-    torch.manual_seed(2)
-    ref = vit_predictor(
-        img_size=(64, 64), patch_size=16, use_mask_tokens=True, embed_dim=768,
-        predictor_embed_dim=384, teacher_embed_dim=None, num_frames=4, tubelet_size=2,
-        depth=12, num_heads=12, num_mask_tokens=8, use_rope=True, uniform_power=False,
-        use_sdpa=False, use_silu=False, wide_silu=True, n_output_distillation=4,
-        return_all_tokens=True, img_temporal_dim_size=1, modality_embedding=True,
-        zero_init_mask_tokens=True, interpolate_rope=True,
-    ).eval()
+def test_predictor_image_modality(video):
+    """The reference spells the video/image switch `mod="image"`; the port uses
+    `mode="img"` and accepts "image" as an alias. A mismatch here silently adds
+    the video modality embedding to an image."""
+    ref_enc, ref_pred, port = make_pair()
+    torch.manual_seed(6)
+    image = torch.randn(2, 3, 1, IMG, IMG)
+    z = ref_enc(image)
+    ctx = torch.arange(0, 10).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(10, 16).unsqueeze(0).expand(2, -1)
+
+    exp_pred, _ = ref_pred(apply_masks(z, [ctx]), [ctx], [tgt], mod="image")
+    for alias in ("img", "image"):
+        got = port.predictor(z, [ctx], [tgt], mode=alias)
+        assert maxdiff(exp_pred, got.last_hidden_state) < 1e-3, alias
+
+    # and the video path must give something different
+    video_out = port.predictor(z, [ctx], [tgt], mode="video")
+    assert not torch.allclose(exp_pred, video_out.last_hidden_state, atol=1e-4)
 
-    converted = convert_reference_predictor_state_dict(ref.state_dict(), 384)
-    missing, unexpected = port.predictor.load_state_dict(converted, strict=False)
-    assert not unexpected, f"unexpected keys: {unexpected}"
-    assert not missing, f"missing keys: {missing}"
 
-    B, N = 2, 32
-    z = torch.randn(B, N, 768 * 4)  # concatenated hierarchical features
-    ctx = torch.arange(0, 20).unsqueeze(0).expand(B, -1)
-    tgt = torch.arange(20, N).unsqueeze(0).expand(B, -1)
+@torch.no_grad()
+def test_full_jepa_forward_matches_reference(video):
+    """Encoder-with-masks followed by the predictor, i.e. the pre-training
+    forward end to end."""
+    ref_enc, ref_pred, port = make_pair()
+    ctx, tgt = _masks()
 
-    from _vjepa21_port_under_test.modeling_vjepa21 import apply_masks
+    z_ref = ref_enc(video, masks=[ctx])                # already gathered
+    exp_pred, exp_ctx = ref_pred(z_ref, [ctx], [tgt], mod="video")
 
-    ref_pred, ref_ctx = ref(apply_masks(z, [ctx]), [ctx], [tgt], mod="video")
-    got = port.predictor(z, [ctx], [tgt], mode="video")
-
-    d_pred = (ref_pred - got.last_hidden_state).abs().max().item()
-    d_ctx = (ref_ctx - got.context_hidden_state).abs().max().item()
-    print(f"[hier predictor target]  max |diff| = {d_pred:.3e}")
-    print(f"[hier predictor context] max |diff| = {d_ctx:.3e}")
-    assert d_pred < 1e-5 and d_ctx < 1e-5
+    out = port(
+        pixel_values_videos=video, masks=[ctx], context_mask=[ctx], target_mask=[tgt]
+    )
+    dp = maxdiff(exp_pred, out.predictor_output.last_hidden_state)
+    dc = maxdiff(exp_ctx, out.predictor_output.context_hidden_state)
+    print(f"\n[full JEPA forward] target max|Δ| = {dp:.3e}   context max|Δ| = {dc:.3e}")
+    assert dp < 1e-3 and dc < 1e-3

tests/test_thesis_robustness.py

@@ -16,6 +16,7 @@ Re-run it against a published checkpoint before trusting a feature store:
 
 from __future__ import annotations
 
+import copy
 import importlib
 import os
 import sys
@@ -47,9 +48,10 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 def tiny_config(**kw):
     base = dict(
         patch_size=16, crop_size=64, hidden_size=96, num_attention_heads=6,
-        num_hidden_layers=12, pred_hidden_size=48, pred_num_attention_heads=6,
-        pred_num_hidden_layers=12, pred_teacher_embed_dim=96,
-        n_output_distillation=1, num_labels=5, attn_implementation="sdpa",
+        num_hidden_layers=12, num_pooler_heads=6, pred_hidden_size=48,
+        pred_num_attention_heads=6, pred_num_hidden_layers=12,
+        pred_teacher_embed_dim=96, n_output_distillation=1, num_labels=5,
+        attn_implementation="sdpa",
     )
     base.update(kw)
     return VJEPA21Config(**base)
@@ -76,6 +78,38 @@ def _rel(a: torch.Tensor, b: torch.Tensor) -> float:
     return ((a - b).abs().mean() / a.abs().mean()).item()
 
 
+def _cos(a: torch.Tensor, b: torch.Tensor) -> float:
+    return torch.nn.functional.cosine_similarity(a.flatten(0, 1), b.flatten(0, 1)).min().item()
+
+
+# --- 0. the fixture must stay pristine ---------------------------------------
+
+
+@torch.no_grad()
+def test_module_fixture_is_not_mutated(model, clip):
+    """Guard for the bug this file used to contain.
+
+    `test_reduced_precision_gap` called `model.to(bfloat16)` on a module-scoped
+    fixture and then `model.to(float32)` to undo it. Casting back does not
+    restore the mantissa: every subsequent test saw bf16-rounded weights stored
+    in float32, and the "fp16 vs fp32" number it reported was really "fp16 vs
+    bf16-rounded", which is far smaller than the truth. Fingerprinting the
+    weights here makes any repeat of that failure loud.
+    """
+    fingerprint = torch.stack(
+        [p.detach().float().flatten()[:8].cpu() for p in list(model.parameters())[:5]]
+    )
+    torch.testing.assert_close(
+        fingerprint,
+        torch.stack(
+            [p.detach().float().flatten()[:8].cpu() for p in list(model.parameters())[:5]]
+        ),
+    )
+    assert all(p.dtype is torch.float32 for p in model.parameters()), (
+        "the shared model fixture is no longer float32; some test cast it in place"
+    )
+
+
 # --- 1. cached features must equal on-the-fly features ------------------------
 
 
@@ -83,12 +117,12 @@ def _rel(a: torch.Tensor, b: torch.Tensor) -> float:
 def test_batch_size_invariance(model, clip):
     """A feature store is precomputed in batches and consumed one clip at a time.
 
-    On CPU this is bit-exact. On GPU it is not: cuBLAS and the fused attention
-    kernels pick tiling and reduction order as a function of the input shape, so
-    the same clip accumulates in a different order at batch 1 and at batch 8. The
-    check is therefore on relative error, and bit-exactness is only required on
-    CPU. `test_determinism_across_calls` covers the same-shape case, which must
-    be exact on both devices.
+    This is not guaranteed to be bit-exact on any device: BLAS and the fused
+    attention kernels pick tiling and reduction order as a function of the input
+    shape, so the same clip can accumulate in a different order at batch 1 and at
+    batch 8. The assertion is therefore on relative error and cosine similarity.
+    `test_determinism_across_calls` covers the same-shape case, which must be
+    exact everywhere.
     """
     full = model(pixel_values_videos=clip, skip_predictor=True).last_hidden_state
     single = torch.cat(
@@ -96,15 +130,9 @@ def test_batch_size_invariance(model, clip):
          for i in range(clip.shape[0])]
     )
     max_abs = (full - single).abs().max().item()
-    rel = _rel(full, single)
-    cos = torch.nn.functional.cosine_similarity(
-        full.flatten(0, 1), single.flatten(0, 1)
-    ).min().item()
+    rel, cos = _rel(full, single), _cos(full, single)
     print(f"\n[batch invariance]   max|Δ| = {max_abs:.3e}   rel = {rel:.3e}   "
-          f"min cos-sim = {cos:.6f}   device = {DEVICE}")
-
-    if DEVICE == "cpu":
-        assert max_abs == 0.0, "batching must not change the result on CPU"
+          f"min cos-sim = {cos:.6f}   device = {DEVICE}   bit-exact = {max_abs == 0.0}")
     assert rel < 1e-4
     assert cos > 0.9999
 
@@ -126,22 +154,31 @@ def test_determinism_across_calls(model, clip):
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
 def test_reduced_precision_gap(model, clip, dtype):
     """Training runs in bf16 or fp16. Measure the gap against the fp32 reference
-    once and record it, rather than assuming it is negligible."""
+    once and record it, rather than assuming it is negligible.
+
+    The low-precision copy is a `deepcopy`: `.to(dtype)` is in-place on the
+    module, and casting back to float32 afterwards does *not* recover the
+    discarded mantissa bits. Mutating the shared fixture would corrupt every
+    later test and, worse, make the second parametrisation measure the gap
+    against an already-degraded reference.
+    """
     if dtype is torch.float16 and DEVICE == "cpu":
         pytest.skip("fp16 matmul is not supported on CPU")
 
     ref = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
+    assert ref.dtype is torch.float32, "the reference forward is not float32"
 
+    low = copy.deepcopy(model).to(dtype)
     try:
-        low = model.to(dtype)
         got = low(pixel_values_videos=clip[:2].to(dtype), skip_predictor=True).last_hidden_state
+        assert got.dtype is dtype, f"output dtype drifted to {got.dtype}"
+        got = got.float()
     finally:
-        model.to(torch.float32)
+        del low
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
 
-    assert got.dtype is dtype, f"output dtype drifted to {got.dtype}"
-    got = got.float()
-    rel = _rel(ref, got)
-    cos = torch.nn.functional.cosine_similarity(ref.flatten(0, 1), got.flatten(0, 1)).min().item()
+    rel, cos = _rel(ref, got), _cos(ref, got)
     print(f"[{str(dtype).split('.')[-1]:>9}]        rel = {rel:.3e}   min cos-sim = {cos:.6f}")
     assert torch.isfinite(got).all()
     assert cos > 0.99
@@ -152,31 +189,32 @@ def test_feature_store_fp16_roundtrip(model, clip):
     """Storing fp32 features as fp16 halves the cache. Quantify what it costs."""
     ref = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
     back = ref.half().float()
-    cos = torch.nn.functional.cosine_similarity(ref.flatten(0, 1), back.flatten(0, 1)).min().item()
-    print(f"[fp16 storage]       rel = {_rel(ref, back):.3e}   min cos-sim = {cos:.6f}")
-    assert cos > 0.9999
+    print(f"[fp16 storage]       rel = {_rel(ref, back):.3e}   min cos-sim = {_cos(ref, back):.6f}")
+    assert _cos(ref, back) > 0.9999
 
 
 # --- 3. the central PEFT claim of the thesis ----------------------------------
-# These three run on a small randomly initialised model regardless of VJEPA21_CKPT:
+# These run on a small randomly initialised model regardless of VJEPA21_CKPT:
 # exact merging is a property of the reparameterisation, not of the weights, and
-# holding a published checkpoint plus a second classification copy does not fit in
-# a notebook runtime for the larger variants. Their output is identical on every
-# repository by construction.
+# holding a published checkpoint plus a second classification copy does not fit
+# in a notebook runtime for the larger variants.
+
+TARGET_MODULES = r".*vjepa21\.encoder\.layer\.\d+\.attention\.(query|key|value|proj)$"
 
 
-def _lora_model(config, r=8):
+def _peft_model(config, r=8, use_dora=False):
     from peft import LoraConfig, get_peft_model
 
     torch.manual_seed(0)
     base = VJEPA21ForVideoClassification(config).eval()
-    lora = LoraConfig(
-        r=r, lora_alpha=2 * r, lora_dropout=0.0,
-        target_modules=r".*vjepa21\.encoder\.layer\.\d+\.attention\.(query|key|value|proj)$",
-        modules_to_save=["classifier"],
+    adapter = LoraConfig(
+        r=r, lora_alpha=2 * r, lora_dropout=0.0, use_dora=use_dora,
+        target_modules=TARGET_MODULES,
+        modules_to_save=["classifier", "pooler"],
     )
-    peft_model = get_peft_model(base, lora)
-    # give the adapters non-zero B matrices, otherwise the merge is trivially exact
+    peft_model = get_peft_model(base, adapter)
+    # Give the adapters non-zero B matrices, otherwise the merge is trivially
+    # exact: peft initialises lora_B to zero, so the adapter is the identity.
     for name, param in peft_model.named_parameters():
         if "lora_B" in name:
             torch.nn.init.normal_(param, std=0.02)
@@ -184,11 +222,18 @@ def _lora_model(config, r=8):
 
 
 @torch.no_grad()
-def test_lora_merge_is_exact():
+@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
+def test_merge_is_exact(use_dora):
     """`merge_and_unload` must not change the function computed. This underpins
-    the claim that reparameterisation methods have zero inference cost."""
+    the claim that reparameterisation methods have zero inference cost.
+
+    DoRA is the case worth testing: it decomposes the update into direction and
+    magnitude, so merging has to fold the magnitude vector back in as well. The
+    model card claimed DoRA was verified while only LoRA was covered.
+    """
+    pytest.importorskip("peft")
     cfg = tiny_config()
-    peft_model = _lora_model(cfg).eval()
+    peft_model = _peft_model(cfg, use_dora=use_dora).eval()
     x = torch.randn(2, 3, 4, cfg.crop_size, cfg.crop_size)
 
     before = peft_model(pixel_values_videos=x).logits
@@ -196,35 +241,60 @@ def test_lora_merge_is_exact():
     after = merged(pixel_values_videos=x).logits
 
     max_abs = (before - after).abs().max().item()
-    print(f"[LoRA merge, tiny]       max|Δ| = {max_abs:.3e}   rel = {_rel(before, after):.3e}")
+    label = "DoRA" if use_dora else "LoRA"
+    print(f"\n[{label} merge, tiny]   max|Δ| = {max_abs:.3e}   rel = {_rel(before, after):.3e}")
     assert max_abs < 1e-4
 
 
 @torch.no_grad()
-def test_merged_model_has_baseline_parameter_count():
+@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
+def test_merged_model_has_baseline_parameter_count(use_dora):
     """After merging there must be no residual adapter parameters: the merged
-    model is parameter-identical to the untouched baseline."""
+    model is parameter-identical to the untouched baseline. This is the
+    zero-inference-cost claim in its structural form, and it is what prompt-based
+    methods cannot satisfy."""
+    pytest.importorskip("peft")
     cfg = tiny_config()
     baseline = VJEPA21ForVideoClassification(cfg)
-    merged = _lora_model(cfg).merge_and_unload()
+    merged = _peft_model(cfg, use_dora=use_dora).merge_and_unload()
 
     n_base = sum(p.numel() for p in baseline.parameters())
     n_merged = sum(p.numel() for p in merged.parameters())
     leftover = [n for n, _ in merged.named_parameters() if "lora" in n.lower()]
-    print(f"[merged params, tiny]    baseline {n_base:,} vs merged {n_merged:,}   "
+    label = "DoRA" if use_dora else "LoRA"
+    print(f"[{label} merged params]  baseline {n_base:,} vs merged {n_merged:,}   "
           f"residual adapter tensors: {len(leftover)}")
     assert n_merged == n_base
     assert not leftover
 
 
 @torch.no_grad()
-def test_unmerged_adapter_adds_parameters():
+@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
+def test_unmerged_adapter_adds_parameters(use_dora):
     """Control for the test above: before merging the adapters really are extra
     parameters, so the equality afterwards is meaningful."""
+    pytest.importorskip("peft")
     cfg = tiny_config()
-    peft_model = _lora_model(cfg)
+    peft_model = _peft_model(cfg, use_dora=use_dora)
     n_base = sum(p.numel() for p in VJEPA21ForVideoClassification(cfg).parameters())
     n_peft = sum(p.numel() for p in peft_model.parameters())
-    print(f"[unmerged params, tiny]    baseline {n_base:,} vs adapted {n_peft:,} "
+    label = "DoRA" if use_dora else "LoRA"
+    print(f"[{label} unmerged]       baseline {n_base:,} vs adapted {n_peft:,} "
           f"(+{n_peft - n_base:,})")
-    assert n_peft > n_base
+    assert n_peft > n_base
+
+
+def test_target_modules_regex_does_not_touch_the_predictor():
+    """`VJEPA21ForVideoClassification` never runs the predictor, so adapting it
+    would spend trainable parameters on dead weight. Matching on bare names such
+    as ["query", "key", "value"] does exactly that."""
+    pytest.importorskip("peft")
+    import re
+
+    cfg = tiny_config()
+    model = VJEPA21ForVideoClassification(cfg)
+    pattern = re.compile(TARGET_MODULES)
+    matched = [n for n, _ in model.named_modules() if pattern.match(n)]
+    assert matched, "the target_modules regex matches nothing"
+    assert not any(".predictor." in n for n in matched)
+    assert len(matched) == cfg.num_hidden_layers * 4

tests/test_vjepa21.py

@@ -44,6 +44,7 @@ VJEPA21Config = CFG_MOD.VJEPA21Config
 VJEPA21Model = MDL_MOD.VJEPA21Model
 VJEPA21ForVideoClassification = MDL_MOD.VJEPA21ForVideoClassification
 VJEPA21VideoProcessor = VPR_MOD.VJEPA21VideoProcessor
+apply_masks = MDL_MOD.apply_masks
 
 
 def tiny_config(**kw):
@@ -54,6 +55,7 @@ def tiny_config(**kw):
         hidden_size=96,
         num_attention_heads=6,
         num_hidden_layers=12,
+        num_pooler_heads=6,     # 96 // 16 = 6 would also work; pin it for clarity
         pred_hidden_size=48,
         pred_num_attention_heads=6,
         pred_num_hidden_layers=12,
@@ -88,6 +90,18 @@ def test_hierarchical_layer_properties():
     assert cfg.pretrained_grid_size == 16
 
 
+def test_predictor_layer_map_differs_from_the_encoder_at_depth_24():
+    """The reference predictor table is [4, 11, 17, 23] at depth 24 while the
+    encoder table is [5, 11, 17, 23]. Sharing one table is a latent bug."""
+    cfg = VJEPA21Config(
+        hidden_size=96, num_attention_heads=6, num_hidden_layers=24,
+        pred_hidden_size=48, pred_num_attention_heads=6, pred_num_hidden_layers=24,
+        n_output_distillation=4, num_pooler_heads=6,
+    )
+    assert cfg.encoder_hierarchical_layers == [5, 11, 17, 23]
+    assert cfg.predictor_hierarchical_layers == [4, 11, 17, 23]
+
+
 def test_config_rejects_inconsistent_values():
     with pytest.raises(ValueError):
         tiny_config(n_output_distillation=9)
@@ -95,11 +109,24 @@ def test_config_rejects_inconsistent_values():
         tiny_config(hidden_size=100, num_attention_heads=6)
     with pytest.raises(ValueError):
         VJEPA21Config(num_hidden_layers=13)
+    with pytest.raises(ValueError):  # predictor depth not in the reference table
+        tiny_config(pred_num_hidden_layers=48)
+    with pytest.raises(ValueError):  # teacher dim not divisible by the level count
+        tiny_config(
+            num_hidden_layers=40, pred_num_hidden_layers=24,
+            n_output_distillation=3, pred_teacher_embed_dim=100,
+        )
 
 
-def test_num_pooler_heads_defaults_to_attention_heads():
-    assert tiny_config().num_pooler_heads == 6
-    assert tiny_config(num_pooler_heads=16, hidden_size=96).num_pooler_heads == 16
+def test_num_pooler_heads_defaults_to_sixteen():
+    """Every frozen-probe config in `configs/eval_2_1/` uses classifier.num_heads: 16,
+    for all four model sizes. Inheriting `num_attention_heads` would give 22 on
+    ViT-g and 26 on ViT-G."""
+    cfg = VJEPA21Config(hidden_size=1408, num_attention_heads=22, num_hidden_layers=40,
+                        n_output_distillation=4, pred_num_hidden_layers=24)
+    assert cfg.num_pooler_heads == 16
+    assert cfg.num_pooler_layers == 3  # + 1 cross-attention block = num_probe_blocks: 4
+    assert VJEPA21Config(num_pooler_heads=8).num_pooler_heads == 8
 
 
 # --- encoder -----------------------------------------------------------------
@@ -130,6 +157,36 @@ def test_unknown_layout_raises(model):
         model(pixel_values_videos=torch.randn(2, 5, 4, 64, 64), skip_predictor=True)
 
 
+def test_ambiguous_layout_warns(recwarn, capsys):
+    """A 3-frame clip in (B, T, C, H, W) also matches (B, C, T, H, W). Deciding
+    silently is the failure mode; the channels-first reading plus a warning is
+    the contract.
+
+    `warning_once` is memoised per message, so the warning is captured by calling
+    `normalize_video_layout` directly with a fresh cache rather than through the
+    module-scoped model, whose first call may already have spent it.
+    """
+    messages = []
+    original = MDL_MOD._warn_once
+    try:
+        MDL_MOD._warn_once = messages.append
+        x = torch.randn(2, 3, 3, 64, 64)
+        out = MDL_MOD.normalize_video_layout(x, in_chans=3)
+    finally:
+        MDL_MOD._warn_once = original
+
+    assert out.shape == (2, 3, 3, 64, 64)          # channels-first reading wins
+    assert torch.equal(out, x)                     # and it is not transposed
+    assert any("ambiguous video layout" in m for m in messages), messages
+
+
+@torch.no_grad()
+def test_ambiguous_layout_still_runs(model):
+    """The warning must not become an exception: a 3-frame clip is legal input."""
+    out = model(pixel_values_videos=torch.randn(2, 3, 3, 64, 64), skip_predictor=True)
+    assert out.last_hidden_state.shape[0] == 2
+
+
 @torch.no_grad()
 def test_output_hidden_states(model):
     out = model(pixel_values_videos=VIDEO, skip_predictor=True, output_hidden_states=True)
@@ -167,6 +224,19 @@ def test_multilevel_and_hierarchical_outputs(model):
     assert out.hierarchical_hidden_state.shape == (2, 32, 96 * n_distill)
 
 
+@torch.no_grad()
+def test_multilevel_output_order_follows_the_network(model):
+    """Levels come back in network order regardless of how they were requested,
+    so a probe that zips them with `encoder_hierarchical_layers` is correct."""
+    hier = model.config.encoder_hierarchical_layers
+    forward = model(pixel_values_videos=VIDEO, skip_predictor=True,
+                    out_layers=hier).multilevel_hidden_states
+    shuffled = model(pixel_values_videos=VIDEO, skip_predictor=True,
+                     out_layers=list(reversed(hier))).multilevel_hidden_states
+    for a, b in zip(forward, shuffled):
+        assert torch.equal(a, b)
+
+
 def test_out_layers_validation(model):
     with pytest.raises(ValueError, match="hierarchical layers"):
         model(pixel_values_videos=VIDEO, skip_predictor=True, out_layers=[3])
@@ -213,6 +283,64 @@ def test_low_precision_weights(dtype):
     assert torch.isfinite(out.last_hidden_state).all()
 
 
+@torch.no_grad()
+def test_unexpected_kwargs_are_reported_not_swallowed(model, caplog):
+    """`out_layer=[11]` (no s) used to be a silent no-op."""
+    import logging
+
+    with caplog.at_level(logging.WARNING, logger=MDL_MOD.logger.name):
+        model(pixel_values_videos=VIDEO, skip_predictor=True, out_layer=[11])
+    # warning_once may have been consumed by another test in the same process,
+    # so only assert that the call did not raise and produced no multilevel output
+    out = model(pixel_values_videos=VIDEO, skip_predictor=True, out_layer=[11])
+    assert out.multilevel_hidden_states is None
+
+
+# --- masked (JEPA) encoder forward -------------------------------------------
+
+
+@torch.no_grad()
+def test_encoder_masks_reduce_the_sequence(model):
+    """`masks` drops tokens before the layers, as `encoder(clips, masks_enc)` does
+    in the reference; attention then only sees the context."""
+    idx = torch.stack([torch.randperm(32)[:12].sort().values for _ in range(2)])
+    out = model.encoder(VIDEO, masks=[idx])
+    assert out.last_hidden_state.shape == (2, 12, 96)
+
+
+@torch.no_grad()
+def test_encoder_masks_carry_true_positions(model):
+    """Masked tokens must keep their original RoPE ids. Feeding the same tokens
+    with a different index set has to change the output, otherwise positions are
+    being renumbered 0..K-1."""
+    keep = torch.arange(0, 12).unsqueeze(0).expand(2, -1)
+    shifted = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    a = model.encoder(VIDEO, masks=[keep]).last_hidden_state
+    b = model.encoder(VIDEO, masks=[shifted]).last_hidden_state
+    assert not torch.allclose(a, b, atol=1e-4)
+
+
+def test_encoder_rejects_out_of_range_masks(model):
+    with pytest.raises(ValueError, match="out of range"):
+        model.encoder(VIDEO, masks=[torch.tensor([[0, 999]]).expand(2, -1)])
+
+
+@torch.no_grad()
+def test_masked_model_forward_runs_the_predictor(model):
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    out = model(pixel_values_videos=VIDEO, masks=[ctx], context_mask=[ctx], target_mask=[tgt])
+    assert out.last_hidden_state.shape == (2, 20, 96)
+    assert out.predictor_output.last_hidden_state.shape == (2, 12, 96)
+    assert out.masked_hidden_state.shape == (2, 20, 96)
+
+
+def test_masked_model_forward_requires_target_mask(model):
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    with pytest.raises(ValueError, match="target_mask"):
+        model(pixel_values_videos=VIDEO, masks=[ctx])
+
+
 # --- predictor ---------------------------------------------------------------
 
 
@@ -234,12 +362,65 @@ def test_predictor_explicit_masks(model):
     assert out.predictor_output.context_hidden_state.shape == (2, 20, 96)
 
 
+@torch.no_grad()
+def test_predictor_accepts_bare_tensors(model):
+    """The reference wraps a bare tensor in a list; so do we."""
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    a = model(pixel_values_videos=VIDEO, context_mask=ctx, target_mask=tgt)
+    b = model(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
+    assert torch.equal(a.predictor_output.last_hidden_state,
+                       b.predictor_output.last_hidden_state)
+
+
+@torch.no_grad()
+def test_mask_index_selects_the_intended_token():
+    """With zero-initialised mask tokens every index behaves identically, which
+    makes any test that relies on them vacuous. Give them distinct values first."""
+    torch.manual_seed(0)
+    m = VJEPA21Model(tiny_config()).eval()
+    for i, tok in enumerate(m.predictor.embeddings.mask_tokens):
+        torch.nn.init.constant_(tok, 0.1 * (i + 1))
+
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    a = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=1)
+    b = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=2)
+    assert not torch.allclose(a.predictor_output.last_hidden_state,
+                              b.predictor_output.last_hidden_state, atol=1e-5)
+    # the index wraps modulo the number of mask tokens, as in the reference
+    n = m.config.pred_num_mask_tokens
+    c = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=1 + n)
+    assert torch.allclose(a.predictor_output.last_hidden_state,
+                          c.predictor_output.last_hidden_state, atol=1e-6)
+
+
+@torch.no_grad()
+def test_predictor_context_is_masked_flag(model):
+    """Passing an already-gathered encoder output must give the same answer as
+    passing the full sequence with the same context indices."""
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    z = model(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
+    full = model.predictor(z, [ctx], [tgt])
+    pre = model.predictor(apply_masks(z, [ctx]), [ctx], [tgt], context_is_masked=True)
+    assert torch.allclose(full.last_hidden_state, pre.last_hidden_state, atol=1e-6)
+
+
 def test_predictor_rejects_multiple_mask_pairs(model):
     m = torch.arange(0, 16).unsqueeze(0).expand(2, -1)
     with pytest.raises(NotImplementedError):
         model(pixel_values_videos=VIDEO, context_mask=[m, m], target_mask=[m, m])
 
 
+def test_predictor_rejects_inconsistent_masked_context(model):
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    z = model(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
+    with pytest.raises(ValueError, match="context_is_masked"):
+        model.predictor(z, [ctx], [tgt], context_is_masked=True)
+
+
 # --- classification head -----------------------------------------------------
 
 
@@ -253,6 +434,28 @@ def test_classification_loss_and_backward():
     assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in clf.parameters())
 
 
+def test_multi_label_classification_loss():
+    """XD-Violence is multi-label; float targets must route to BCE."""
+    torch.manual_seed(0)
+    cfg = tiny_config()
+    cfg.problem_type = "multi_label_classification"
+    clf = VJEPA21ForVideoClassification(cfg)
+    labels = torch.zeros(2, 5)
+    labels[0, 1] = labels[0, 4] = labels[1, 2] = 1.0
+    out = clf(pixel_values_videos=VIDEO, labels=labels)
+    assert torch.isfinite(out.loss)
+    out.loss.backward()
+
+
+def test_trainer_style_kwargs_do_not_reach_the_encoder():
+    """The Trainer injects `num_items_in_batch`; it must not blow up or be
+    forwarded down to the encoder."""
+    clf = VJEPA21ForVideoClassification(tiny_config()).eval()
+    with torch.no_grad():
+        out = clf(pixel_values_videos=VIDEO, num_items_in_batch=2)
+    assert out.logits.shape == (2, 5)
+
+
 def test_gradient_checkpointing_matches_plain_forward():
     torch.manual_seed(0)
     a = VJEPA21ForVideoClassification(tiny_config())
@@ -267,6 +470,28 @@ def test_gradient_checkpointing_matches_plain_forward():
     )
 
 
+def test_gradient_checkpointing_matches_on_the_predictor_path():
+    """The classification head never runs the predictor, so its checkpointing
+    path is otherwise untested."""
+    torch.manual_seed(0)
+    a = VJEPA21Model(tiny_config())
+    b = VJEPA21Model(tiny_config())
+    b.load_state_dict(a.state_dict())
+    a.train()
+    b.train()
+    b.gradient_checkpointing_enable()
+    assert b.predictor.gradient_checkpointing
+    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
+    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
+    ta = a(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
+    tb = b(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
+    assert torch.allclose(
+        ta.predictor_output.last_hidden_state,
+        tb.predictor_output.last_hidden_state,
+        atol=1e-5,
+    )
+
+
 @torch.no_grad()
 def test_classification_propagates_hidden_states():
     clf = VJEPA21ForVideoClassification(tiny_config()).eval()
@@ -308,6 +533,19 @@ def test_video_processor_crop_size_is_configurable():
     assert vp.size == {"shortest_edge": 256}
 
 
+def test_processor_output_feeds_the_model_without_a_transpose():
+    """The processor emits (B, T, C, H, W); the model must accept it as-is."""
+    vp = VJEPA21VideoProcessor(crop_size=64)
+    clip = [np.random.randint(0, 255, (80, 100, 3), dtype=np.uint8) for _ in range(4)]
+    pv = vp([clip], return_tensors="pt")["pixel_values_videos"]
+    assert pv.shape == (1, 4, 3, 64, 64)
+    m = VJEPA21Model(tiny_config()).eval()
+    with torch.no_grad():
+        assert m(pixel_values_videos=pv, skip_predictor=True).last_hidden_state.shape == (
+            1, 2 * 4 * 4, 96
+        )
+
+
 # --- serialisation round-trip ------------------------------------------------
 
 
@@ -380,6 +618,19 @@ def test_backbone_weights_survive_the_classification_wrapper():
         shutil.rmtree(tmp, ignore_errors=True)
 
 
+def test_saved_config_round_trips_the_pooler_heads():
+    """A regression here would silently change the probe architecture."""
+    tmp = tempfile.mkdtemp()
+    try:
+        cfg = tiny_config(num_pooler_heads=6)
+        VJEPA21Model(cfg).save_pretrained(tmp)
+        reloaded = VJEPA21Config.from_pretrained(tmp)
+        assert reloaded.num_pooler_heads == 6
+        assert reloaded.num_pooler_layers == cfg.num_pooler_layers
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
 # --- multi-level distillation variants (ViT-g / ViT-G style) ------------------
 
 
@@ -387,6 +638,7 @@ def giant_like_config(**kw):
     base = dict(
         patch_size=16, crop_size=64, hidden_size=96, num_attention_heads=6,
         num_hidden_layers=40, mlp_ratio=48 / 11, n_output_distillation=4,
+        num_pooler_heads=6,
         pred_hidden_size=48, pred_num_attention_heads=6, pred_num_hidden_layers=24,
         pred_num_mask_tokens=8, pred_teacher_embed_dim=None, pred_return_all_tokens=True,
     )
@@ -406,6 +658,8 @@ def test_predictor_consumes_hierarchical_features_when_required():
     out = model(pixel_values_videos=x)  # would raise a shape error before the fix
     assert out.predictor_output.last_hidden_state.shape == (1, 32, 96 * 4)
     assert out.predictor_output.context_hidden_state.shape == (1, 32, 96 * 4)
+    # masked_hidden_state carries the hierarchical width in this regime
+    assert out.masked_hidden_state.shape == (1, 32, 96 * 4)
 
 
 @torch.no_grad()

video_processing_vjepa21.py

@@ -6,9 +6,21 @@ Reproduces the evaluation-time transform of the reference implementation
     resize the short side to `crop_size` -> square crop of `crop_size`
     -> scale to [0, 1] -> normalise with ImageNet statistics
 
-Note that this differs from the stock `VJEPA2VideoProcessor`, which resizes the
-short side to `crop_size * 256 / 224` before cropping. The reference V-JEPA 2.1
-evaluations resize to exactly the crop size, so a dedicated class is required.
+Two things are worth stating precisely.
+
+*Which reference transform this is.* `make_transforms` picks `EvalVideoTransform`
+when `num_views_per_clip > 1` and `VideoTransform(training=False)` otherwise, and
+only the first resizes the short side to exactly `crop_size` — the second uses
+`crop_size * 256 / 224`, like the stock `VJEPA2VideoProcessor`. Every V-JEPA 2.1
+frozen-probe config in `configs/eval_2_1/` sets `num_views_per_segment: 3`, so
+`EvalVideoTransform` is the transform the 2.1 evaluations actually use and
+resizing to exactly `crop_size` is correct here.
+
+*What is not reproduced.* `EvalVideoTransform` takes `num_views_per_clip` crops
+sliding along the long axis and the evaluation loop averages their softmax
+outputs. This processor takes a single centre crop. Multi-view aggregation is
+the caller's job; `video_io.aggregate_predictions` implements the reference
+combination rule.
 
 Compatible with transformers 4.x and 5.x. In particular it does not declare a
 `VideosKwargs` subclass as `valid_kwargs`: on transformers 5 an empty TypedDict
@@ -19,7 +31,12 @@ is correct on both majors.
 
 import torch
 
-from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling, SizeDict
+from transformers.image_utils import (
+    IMAGENET_DEFAULT_MEAN,
+    IMAGENET_DEFAULT_STD,
+    PILImageResampling,
+    SizeDict,
+)
 from transformers.processing_utils import Unpack, VideosKwargs
 from transformers.video_processing_utils import BaseVideoProcessor