modeling_vjepa21.py

"""V-JEPA 2.1 model implementation for HuggingFace Transformers.

Faithful port of the reference implementation in
``facebookresearch/vjepa2`` under ``app/vjepa_2_1/models/``.

Key differences from V-JEPA 2:
- Multi-modality: separate patch embedding for images (tubelet_size=1) + modality embeddings
- Hierarchical output: intermediate layer features with per-layer norms
- Interpolatable RoPE: variable input resolution support
- 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

from collections.abc import Callable
from dataclasses import dataclass
from typing import Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from transformers.activations import ACT2FN
from transformers.modeling_outputs import ImageClassifierOutput
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
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"}
)


# ---------------------------------------------------------------------------
# Output dataclasses
# ---------------------------------------------------------------------------


@dataclass
class VJEPA21EncoderOutput(ModelOutput):
    """Encoder output.

    Attributes:
        last_hidden_state: Final layer output, normalised with the last
            hierarchical norm, `(B, N, hidden_size)`.
        hierarchical_hidden_state: Distillation levels concatenated along the
            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, 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
            `output_attentions=True`.
    """

    last_hidden_state: Optional[torch.FloatTensor] = None
    hierarchical_hidden_state: Optional[torch.FloatTensor] = None
    multilevel_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[tuple[torch.FloatTensor, ...]] = None


@dataclass
class VJEPA21PredictorOutput(ModelOutput):
    """Predictor output.

    Attributes:
        last_hidden_state: Predicted target tokens `(B, N_target, proj_dim)`.
        context_hidden_state: Predicted context tokens, when the model was
            configured with `pred_return_all_tokens=True`.
        hidden_states: Raw outputs of every predictor layer.
        attentions: Attention probabilities of every predictor layer.
    """

    last_hidden_state: Optional[torch.FloatTensor] = None
    context_hidden_state: Optional[torch.FloatTensor] = None
    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[tuple[torch.FloatTensor, ...]] = None


@dataclass
class VJEPA21ModelOutput(ModelOutput):
    """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
    multilevel_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
    masked_hidden_state: Optional[torch.FloatTensor] = None
    predictor_output: Optional[VJEPA21PredictorOutput] = None
    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[tuple[torch.FloatTensor, ...]] = None


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------


def apply_masks(tensor: torch.Tensor, masks: list[torch.Tensor]) -> torch.Tensor:
    """Gather tokens at mask indices.

    Args:
        tensor: `(B, N, D)` tensor.
        masks: List of `(B, K)` index tensors.

    Returns:
        `(len(masks)*B, K, D)` gathered tensor.
    """
    parts = []
    for mask in masks:
        mask = mask.to(tensor.device)
        idx = mask.unsqueeze(-1).expand(-1, -1, tensor.size(-1))
        parts.append(torch.gather(tensor, dim=1, index=idx))
    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)`.

    Accepted layouts, resolved in this order:
        - `(B, C, T, H, W)` channels-first (native)
        - `(B, T, C, H, W)` the layout produced by HF video processors
        - `(B, T, H, W, C)` channels-last
        - `(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 — 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:
        # (B, C, H, W) -> (B, C, 1, H, W)
        if x.shape[1] != in_chans:
            raise ValueError(
                f"Expected a 4D tensor shaped (B, {in_chans}, H, W), got {tuple(x.shape)}."
            )
        return x.unsqueeze(2)

    if x.ndim != 5:
        raise ValueError(
            "pixel_values_videos must be a 4D or 5D tensor, got "
            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)
        return x.permute(0, 2, 1, 3, 4)
    if x.shape[-1] == in_chans:  # (B, T, H, W, C)
        return x.permute(0, 4, 1, 2, 3)

    raise ValueError(
        f"Could not locate a channel axis of size {in_chans} in a tensor of shape "
        f"{tuple(x.shape)}. Supported layouts are (B, C, T, H, W), (B, T, C, H, W) "
        "and (B, T, H, W, C)."
    )


def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
    if drop_prob == 0.0 or not training:
        return x
    keep = 1.0 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
    mask = keep + torch.rand(shape, dtype=x.dtype, device=x.device)
    mask.floor_()
    return x.div(keep) * mask


class VJEPA21DropPath(nn.Module):
    def __init__(self, p: Optional[float] = None):
        super().__init__()
        self.p = p

    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,
    pos: torch.Tensor,
    n_registers: int = 0,
    has_cls_first: bool = False,
) -> torch.Tensor:
    """Apply rotary position embeddings with register/CLS token handling.

    Args:
        x: `(B, H, N, D)` query or key tensor.
        pos: Position ids broadcastable to `(..., N_ctx)`.
        n_registers: Number of register tokens at end of sequence (not rotated).
        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

    x_cls = x[..., :n_cls, :] if n_cls else None
    x_ctx = x[..., start_ctx:end_ctx, :]
    x_reg = x[..., end_ctx:, :] if n_registers > 0 else None

    # Position ids are computed over the context tokens only. If they were built
    # over the full sequence, trim them so CLS/registers stay unrotated.
    if pos.shape[-1] == N and (n_cls or n_registers):
        pos = pos[..., start_ctx:end_ctx]

    # RoPE frequencies are computed in float32 and cast back to the input dtype.
    # The reference implementation builds them in `x.dtype`, which is equivalent in
    # float32 but promotes q/k to float32 when the weights are held in bf16/fp16 —
    # the fused attention kernels then reject the mismatch against v. Meta's training
    # code never hits this because it runs under `torch.autocast` rather than casting
    # the weights. Computing in float32 is also the numerically safer choice.
    omega = torch.arange(D // 2, dtype=torch.float32, device=x.device)
    omega = omega / (D / 2.0)
    omega = 1.0 / (10000.0**omega)

    freq = torch.einsum("..., f -> ... f", pos.to(torch.float32), omega)
    emb_sin = freq.sin().repeat_interleave(2, dim=-1).to(x.dtype)
    emb_cos = freq.cos().repeat_interleave(2, dim=-1).to(x.dtype)

    y = x_ctx.unflatten(-1, (-1, 2))
    y1, y2 = y.unbind(dim=-1)
    y = torch.stack((-y2, y1), dim=-1).flatten(-2)
    out_ctx = x_ctx * emb_cos + y * emb_sin

    parts = []
    if x_cls is not None:
        parts.append(x_cls)
    parts.append(out_ctx)
    if x_reg is not None:
        parts.append(x_reg)
    return torch.cat(parts, dim=-2)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs,
):
    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value).transpose(1, 2).contiguous()
    return attn_output, attn_weights


def resolve_attention_interface(config, output_attentions: bool = False) -> Callable:
    """Return the attention kernel for ``config._attn_implementation``.

    Works on transformers 4.x, where ``ALL_ATTENTION_FUNCTIONS`` is a plain
    mapping, and on 5.x, where it exposes ``get_interface``. When attention
    probabilities are requested we fall back to the eager kernel, since the
    fused kernels do not materialise them.
    """
    if output_attentions:
        return eager_attention_forward

    impl = getattr(config, "_attn_implementation", None) or "eager"
    if impl == "eager":
        return eager_attention_forward

    getter = getattr(ALL_ATTENTION_FUNCTIONS, "get_interface", None)
    if getter is not None:  # transformers >= 5
        return getter(impl, eager_attention_forward)
    return ALL_ATTENTION_FUNCTIONS.get(impl, eager_attention_forward)  # transformers 4.x


# ---------------------------------------------------------------------------
# Patch Embeddings
# ---------------------------------------------------------------------------


class VJEPA21PatchEmbeddings3D(nn.Module):
    """3D patch embedding via Conv3d."""

    def __init__(self, config: VJEPA21Config, tubelet_size: Optional[int] = None):
        super().__init__()
        ts = tubelet_size if tubelet_size is not None else config.tubelet_size
        ps = config.patch_size
        self.proj = nn.Conv3d(
            config.in_chans,
            config.hidden_size,
            kernel_size=(ts, ps, ps),
            stride=(ts, ps, ps),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, C, T, H, W)
        return self.proj(x).flatten(2).transpose(1, 2)


class VJEPA21Embeddings(nn.Module):
    """Patch embeddings with modality-aware processing."""

    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config

        # Video patch embedding (tubelet_size from config)
        self.patch_embeddings = VJEPA21PatchEmbeddings3D(config)

        # Image patch embedding (tubelet_size=1) if img_temporal_dim_size is set
        self.patch_embeddings_img = None
        if config.img_temporal_dim_size is not None:
            self.patch_embeddings_img = VJEPA21PatchEmbeddings3D(config, tubelet_size=1)

        # Modality embeddings
        self.img_mod_embed = None
        self.video_mod_embed = None
        if config.modality_embedding:
            self.img_mod_embed = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
            self.video_mod_embed = nn.Parameter(torch.zeros(1, 1, config.hidden_size))

    def forward(self, pixel_values_videos: torch.Tensor) -> tuple[torch.Tensor, str]:
        """
        Args:
            pixel_values_videos: `(B, C, T, H, W)`, already layout-normalised.

        Returns:
            embeddings: `(B, N, hidden_size)`.
            mode: "img" or "video".
        """
        target_dtype = self.patch_embeddings.proj.weight.dtype
        pixel_values_videos = pixel_values_videos.to(dtype=target_dtype)

        T = pixel_values_videos.shape[2]
        is_image = (
            self.config.img_temporal_dim_size is not None
            and T == self.config.img_temporal_dim_size
        )

        if is_image and self.patch_embeddings_img is not None:
            embeddings = self.patch_embeddings_img(pixel_values_videos)
            mode = "img"
        else:
            # Ensure at least tubelet_size frames
            if T < self.config.tubelet_size:
                pixel_values_videos = pixel_values_videos.repeat(
                    1, 1, self.config.tubelet_size, 1, 1
                )
            embeddings = self.patch_embeddings(pixel_values_videos)
            mode = "video"

        if self.img_mod_embed is not None:
            if mode == "img":
                embeddings = embeddings + self.img_mod_embed
            else:
                embeddings = embeddings + self.video_mod_embed

        return embeddings, mode


# ---------------------------------------------------------------------------
# Attention
# ---------------------------------------------------------------------------


class VJEPA21RopeAttention(nn.Module):
    """RoPE-based multi-head attention with interpolation and register support."""

    def __init__(self, config: VJEPA21Config, hidden_size: int, num_attention_heads: int):
        super().__init__()
        self.config = config
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.attention_head_size = hidden_size // num_attention_heads
        self.all_head_size = num_attention_heads * self.attention_head_size

        self.query = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias)
        self.key = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias)
        self.value = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias)
        self.proj = nn.Linear(hidden_size, hidden_size)

        self.dropout_prob = config.attention_probs_dropout_prob
        self.scaling = self.attention_head_size**-0.5
        self.is_causal = False

        # RoPE dimension split: depth, height, width
        self.d_dim = int(2 * ((self.attention_head_size // 3) // 2))
        self.h_dim = int(2 * ((self.attention_head_size // 3) // 2))
        self.w_dim = int(2 * ((self.attention_head_size // 3) // 2))

        self.grid_size = config.crop_size // config.patch_size
        self.n_registers = config.n_registers
        self.has_cls_first = config.has_cls_first
        self.interpolate_rope = config.interpolate_rope
        self.pretrained_grid_size = config.pretrained_grid_size

    def _separate_positions(
        self,
        ids: torch.Tensor,
        H_patches: Optional[int] = None,
        W_patches: Optional[int] = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Decompose flat token ids into (depth, height, width) components."""
        hp = H_patches if H_patches is not None else self.grid_size
        wp = W_patches if W_patches is not None else self.grid_size
        tokens_per_frame = hp * wp
        frame_ids = ids // tokens_per_frame
        remainder = ids - tokens_per_frame * frame_ids
        height_ids = remainder // wp
        width_ids = remainder - wp * height_ids
        return frame_ids.float(), height_ids.float(), width_ids.float()

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_mask: Optional[torch.Tensor] = None,
        T: Optional[int] = None,
        H_patches: Optional[int] = None,
        W_patches: Optional[int] = None,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        B, N, C = hidden_states.shape

        q = (
            self.query(hidden_states)
            .view(B, N, self.num_attention_heads, self.attention_head_size)
            .transpose(1, 2)
        )
        k = (
            self.key(hidden_states)
            .view(B, N, self.num_attention_heads, self.attention_head_size)
            .transpose(1, 2)
        )
        v = (
            self.value(hidden_states)
            .view(B, N, self.num_attention_heads, self.attention_head_size)
            .transpose(1, 2)
        )

        # Compute position ids
        if position_mask is not None:
            ids = position_mask.unsqueeze(1).repeat(1, self.num_attention_heads, 1)
        else:
            ids = torch.arange(N, device=hidden_states.device)

        d_mask, h_mask, w_mask = self._separate_positions(ids, H_patches, W_patches)

        # Interpolate RoPE for variable resolution.
        # Mirrors app/vjepa_2_1/models/utils/modules.py (`interpolate_rope`).
        if self.interpolate_rope:
            hp = H_patches if H_patches is not None else self.grid_size
            wp = W_patches if W_patches is not None else self.grid_size
            h_mask = h_mask * (self.pretrained_grid_size - 1) / max(hp - 1, 1)
            w_mask = w_mask * (self.pretrained_grid_size - 1) / max(wp - 1, 1)

        # Apply RoPE to each dimension slice
        s = 0
        qd = rotate_queries_or_keys(
            q[..., s : s + self.d_dim], d_mask, self.n_registers, self.has_cls_first
        )
        kd = rotate_queries_or_keys(
            k[..., s : s + self.d_dim], d_mask, self.n_registers, self.has_cls_first
        )
        s += self.d_dim
        qh = rotate_queries_or_keys(
            q[..., s : s + self.h_dim], h_mask, self.n_registers, self.has_cls_first
        )
        kh = rotate_queries_or_keys(
            k[..., s : s + self.h_dim], h_mask, self.n_registers, self.has_cls_first
        )
        s += self.h_dim
        qw = rotate_queries_or_keys(
            q[..., s : s + self.w_dim], w_mask, self.n_registers, self.has_cls_first
        )
        kw = rotate_queries_or_keys(
            k[..., s : s + self.w_dim], w_mask, self.n_registers, self.has_cls_first
        )
        s += self.w_dim

        if s < self.attention_head_size:
            q = torch.cat([qd, qh, qw, q[..., s:]], dim=-1)
            k = torch.cat([kd, kh, kw, k[..., s:]], dim=-1)
        else:
            q = torch.cat([qd, qh, qw], dim=-1)
            k = torch.cat([kd, kh, kw], dim=-1)

        attention_interface = resolve_attention_interface(self.config, output_attentions)

        context_layer, attn_weights = attention_interface(
            self,
            q,
            k,
            v,
            None,
            is_causal=self.is_causal,
            scaling=self.scaling,
            dropout=0.0 if not self.training else self.dropout_prob,
        )

        output = self.proj(context_layer.reshape(B, N, self.all_head_size))
        return output, attn_weights


# ---------------------------------------------------------------------------
# MLP
# ---------------------------------------------------------------------------


class VJEPA21MLP(nn.Module):
    """Standard GELU MLP."""

    def __init__(self, config: VJEPA21Config, hidden_size: int, mlp_ratio: float = 4.0):
        super().__init__()
        hidden_features = int(hidden_size * mlp_ratio)
        self.fc1 = nn.Linear(hidden_size, hidden_features)
        self.act = ACT2FN[config.hidden_act]
        self.fc2 = nn.Linear(hidden_features, hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.act(self.fc1(x)))


class VJEPA21SwiGLUMLP(nn.Module):
    """SwiGLU FFN as used in V-JEPA 2.1."""

    def __init__(self, config: VJEPA21Config, hidden_size: int, mlp_ratio: float = 4.0):
        super().__init__()
        hidden_features = int(hidden_size * mlp_ratio)
        if config.wide_silu:
            swiglu_hidden = int(2 * hidden_features / 3)
            align_as = 8
            swiglu_hidden = (swiglu_hidden + align_as - 1) // align_as * align_as
        else:
            swiglu_hidden = hidden_features
        self.fc1 = nn.Linear(hidden_size, swiglu_hidden)
        self.fc2 = nn.Linear(hidden_size, swiglu_hidden)
        self.fc3 = nn.Linear(swiglu_hidden, hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))


def _make_mlp(config: VJEPA21Config, hidden_size: int, mlp_ratio: float) -> nn.Module:
    if config.hidden_act == "silu":
        return VJEPA21SwiGLUMLP(config, hidden_size, mlp_ratio)
    return VJEPA21MLP(config, hidden_size, mlp_ratio)


# ---------------------------------------------------------------------------
# Transformer Layer
# ---------------------------------------------------------------------------


class VJEPA21Layer(nn.Module):
    """Single transformer block: LN -> Attention -> DropPath + Residual -> LN -> MLP -> DropPath + Residual."""

    def __init__(
        self,
        config: VJEPA21Config,
        hidden_size: int,
        num_attention_heads: int,
        mlp_ratio: float,
        drop_path_rate: float = 0.0,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
        self.attention = VJEPA21RopeAttention(config, hidden_size, num_attention_heads)
        self.drop_path = (
            VJEPA21DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
        )
        self.norm2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
        self.mlp = _make_mlp(config, hidden_size, mlp_ratio)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_mask: Optional[torch.Tensor] = None,
        T: Optional[int] = None,
        H_patches: Optional[int] = None,
        W_patches: Optional[int] = None,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        residual = hidden_states
        h = self.norm1(hidden_states)
        attn_out, attn_weights = self.attention(
            h, position_mask, T, H_patches, W_patches, output_attentions
        )
        hidden_states = residual + self.drop_path(attn_out)

        residual = hidden_states
        hidden_states = residual + self.drop_path(self.mlp(self.norm2(hidden_states)))
        return hidden_states, attn_weights


# ---------------------------------------------------------------------------
# Encoder
# ---------------------------------------------------------------------------


class VJEPA21Encoder(nn.Module):
    """V-JEPA 2.1 encoder with hierarchical and multi-level output support."""

    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config
        self.embeddings = VJEPA21Embeddings(config)

        dpr = [
            config.drop_path_rate * i / max(config.num_hidden_layers - 1, 1)
            for i in range(config.num_hidden_layers)
        ]
        self.layer = nn.ModuleList(
            [
                VJEPA21Layer(
                    config,
                    hidden_size=config.hidden_size,
                    num_attention_heads=config.num_attention_heads,
                    mlp_ratio=config.mlp_ratio,
                    drop_path_rate=dpr[i],
                )
                for i in range(config.num_hidden_layers)
            ]
        )

        # Per-layer norms for hierarchical outputs. One norm per hierarchical
        # level, exactly as in the reference VisionTransformer.
        hier_layers = config.encoder_hierarchical_layers
        self.norms_block = nn.ModuleList(
            [nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) for _ in hier_layers]
        )
        self._hier_layers = hier_layers
        self._distill_layers = config.encoder_distillation_layers

        self.gradient_checkpointing = False

    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
        )

        embeddings, _ = self.embeddings(pixel_values_videos)

        B, C, T_raw, H, W = pixel_values_videos.shape
        is_image = (
            self.config.img_temporal_dim_size is not None
            and T_raw == self.config.img_temporal_dim_size
        )
        T_patches = T_raw if is_image else max(T_raw // self.config.tubelet_size, 1)
        H_patches = H // self.config.patch_size
        W_patches = W // self.config.patch_size

        if out_layers is not None:
            unknown = [i for i in out_layers if i not in self._hier_layers]
            if unknown:
                raise ValueError(
                    f"out_layers={out_layers} contains indices {unknown} that are not "
                    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: 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:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_out = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    position_mask,
                    T_patches,
                    H_patches,
                    W_patches,
                    output_attentions,
                )
            else:
                layer_out = layer_module(
                    hidden_states,
                    position_mask=position_mask,
                    T=T_patches,
                    H_patches=H_patches,
                    W_patches=W_patches,
                    output_attentions=output_attentions,
                )
            hidden_states = layer_out[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_out[1],)

            if out_layers is not None and i in out_layers:
                idx = self._hier_layers.index(i)
                multilevel_outputs.append(self.norms_block[idx](hidden_states))

            if i in self._distill_layers:
                idx = self._hier_layers.index(i)
                hier_outputs.append(self.norms_block[idx](hidden_states))

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        # Reference implementation: `x = self.norms_block[-1](x)` on the final layer.
        last_hidden_state = self.norms_block[-1](hidden_states)

        hierarchical_hidden_state = None
        if return_hierarchical and hier_outputs:
            hierarchical_hidden_state = torch.cat(hier_outputs, dim=2)

        return VJEPA21EncoderOutput(
            last_hidden_state=last_hidden_state,
            hierarchical_hidden_state=hierarchical_hidden_state,
            multilevel_hidden_states=tuple(multilevel_outputs) if out_layers else None,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


# ---------------------------------------------------------------------------
# Predictor
# ---------------------------------------------------------------------------


class VJEPA21PredictorEmbeddings(nn.Module):
    """Predictor embeddings with hierarchical input fusion."""

    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config
        n_hier = len(config.predictor_hierarchical_layers)
        if n_hier <= 1:
            self.predictor_embed = nn.Linear(config.hidden_size, config.pred_hidden_size)
        else:
            act = nn.SiLU if config.hidden_act == "silu" else nn.GELU
            self.predictor_embed = nn.Sequential(
                nn.Linear(config.hidden_size * n_hier, config.hidden_size),
                act(),
                nn.Linear(config.hidden_size, config.pred_hidden_size),
            )

        self.num_mask_tokens = config.pred_num_mask_tokens
        self.mask_tokens = nn.ParameterList(
            [
                nn.Parameter(torch.zeros(1, 1, config.pred_hidden_size))
                for _ in range(self.num_mask_tokens)
            ]
        )

        self.img_mod_embed = None
        self.video_mod_embed = None
        if config.img_temporal_dim_size is not None and config.modality_embedding:
            self.img_mod_embed = nn.Parameter(torch.zeros(1, 1, config.pred_hidden_size))
            self.video_mod_embed = nn.Parameter(torch.zeros(1, 1, config.pred_hidden_size))

    def forward(
        self,
        hidden_states: torch.Tensor,
        context_mask: list[torch.Tensor],
        target_mask: list[torch.Tensor],
        mask_index: int = 1,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # `hidden_states` already carries len(context_mask) * B rows.
        batch_size = hidden_states.size(0) // len(context_mask)

        context = self.predictor_embed(hidden_states)

        mask_index = mask_index % self.num_mask_tokens
        pred_tokens = self.mask_tokens[mask_index].repeat(
            batch_size, self._max_patches(target_mask), 1
        )
        pred_tokens = apply_masks(pred_tokens, target_mask)

        context = context.repeat(len(context_mask), 1, 1)
        embeddings = torch.cat([context, pred_tokens], dim=1)

        cm = torch.cat(context_mask, dim=0)
        tm = torch.cat(target_mask, dim=0)
        masks = torch.cat([cm, tm], dim=1)
        return embeddings, masks

    @staticmethod
    def _max_patches(masks: list[torch.Tensor]) -> int:
        return int(max(m.max().item() for m in masks)) + 1


class VJEPA21Predictor(nn.Module):
    """V-JEPA 2.1 predictor with hierarchical input and context projection."""

    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config
        self.embeddings = VJEPA21PredictorEmbeddings(config)

        dpr = [
            config.drop_path_rate * i / max(config.pred_num_hidden_layers - 1, 1)
            for i in range(config.pred_num_hidden_layers)
        ]
        self.layer = nn.ModuleList(
            [
                VJEPA21Layer(
                    config,
                    hidden_size=config.pred_hidden_size,
                    num_attention_heads=config.pred_num_attention_heads,
                    mlp_ratio=config.pred_mlp_ratio,
                    drop_path_rate=dpr[i],
                )
                for i in range(config.pred_num_hidden_layers)
            ]
        )
        self.layernorm = nn.LayerNorm(config.pred_hidden_size, eps=config.layer_norm_eps)

        n_hier = len(config.predictor_hierarchical_layers)
        if config.pred_teacher_embed_dim is not None:
            out_dim = config.pred_teacher_embed_dim // n_hier
        else:
            out_dim = config.hidden_size
        proj_out_dim = n_hier * out_dim

        self.proj = nn.Linear(config.pred_hidden_size, proj_out_dim)
        self.proj_context = None
        if config.pred_return_all_tokens:
            self.proj_context = nn.Linear(config.pred_hidden_size, proj_out_dim)

        self.gradient_checkpointing = False

    def forward(
        self,
        encoder_hidden_states: 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 "
                f"{len(context_mask)} and {len(target_mask)}."
            )
        if len(context_mask) > 1:
            raise NotImplementedError(
                "The predictor currently supports a single (context_mask, target_mask) "
                "pair. Call it once per mask pair instead."
            )

        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, mask_index=mask_index
        )

        # Sort tokens by position so RoPE sees monotonically increasing ids
        argsort = torch.argsort(position_masks, dim=1)
        idx_expand = argsort.unsqueeze(-1).expand(-1, -1, hidden_states.size(-1))
        hidden_states = torch.gather(hidden_states, 1, idx_expand.to(hidden_states.device))
        position_masks = torch.gather(position_masks, 1, argsort.to(position_masks.device))

        if self.embeddings.img_mod_embed is not None:
            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: 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:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_out = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    position_masks,
                    None,
                    None,
                    None,
                    output_attentions,
                )
            else:
                layer_out = layer_module(
                    hidden_states,
                    position_mask=position_masks,
                    output_attentions=output_attentions,
                )
            hidden_states = layer_out[0]
            if output_attentions:
                all_attentions = all_attentions + (layer_out[1],)

        hidden_states = self.layernorm(hidden_states)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        # Unsort
        reverse = torch.argsort(argsort, dim=1)
        rev_expand = reverse.unsqueeze(-1).expand(-1, -1, hidden_states.size(-1))
        hidden_states = torch.gather(hidden_states, 1, rev_expand.to(hidden_states.device))

        pred = self.proj(hidden_states[:, N_ctxt:])
        ctx = None
        if self.config.pred_return_all_tokens:
            ctx = self.proj_context(hidden_states[:, :N_ctxt])

        return VJEPA21PredictorOutput(
            last_hidden_state=pred,
            context_hidden_state=ctx,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


# ---------------------------------------------------------------------------
# Attentive Pooler (for downstream tasks)
# ---------------------------------------------------------------------------


class VJEPA21PoolerSelfAttention(nn.Module):
    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_pooler_heads
        self.head_dim = self.embed_dim // self.num_heads
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_probs_dropout_prob
        self.is_causal = False

        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

    def forward(
        self, hidden_states: torch.Tensor, output_attentions: bool = False
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        B, N, C = hidden_states.shape
        q = self.q_proj(hidden_states).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(hidden_states).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(hidden_states).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)

        attention_interface = resolve_attention_interface(self.config, output_attentions)
        attn_output, attn_weights = attention_interface(
            self,
            q,
            k,
            v,
            None,
            is_causal=False,
            scaling=self.scale,
            dropout=0.0 if not self.training else self.dropout,
        )
        return self.out_proj(attn_output.reshape(B, N, C)), attn_weights


class VJEPA21PoolerCrossAttention(nn.Module):
    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_pooler_heads
        self.head_dim = self.embed_dim // self.num_heads
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_probs_dropout_prob
        self.is_causal = False

        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)

    def forward(
        self, queries: torch.Tensor, kv: torch.Tensor, output_attentions: bool = False
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        B, Nq, C = queries.shape
        Nkv = kv.shape[1]
        q = self.q_proj(queries).view(B, Nq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(kv).view(B, Nkv, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(kv).view(B, Nkv, self.num_heads, self.head_dim).transpose(1, 2)

        attention_interface = resolve_attention_interface(self.config, output_attentions)
        attn_output, attn_weights = attention_interface(
            self,
            q,
            k,
            v,
            None,
            is_causal=False,
            scaling=self.scale,
            dropout=0.0 if not self.training else self.dropout,
        )
        return attn_output.reshape(B, Nq, C), attn_weights


class VJEPA21PoolerSelfAttentionLayer(nn.Module):
    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.self_attn = VJEPA21PoolerSelfAttention(config)
        self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp = VJEPA21MLP(config, hidden_size=config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states, _ = self.self_attn(self.layer_norm1(hidden_states))
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = residual + self.mlp(self.layer_norm2(hidden_states))
        return hidden_states


class VJEPA21PoolerCrossAttentionLayer(nn.Module):
    def __init__(self, config: VJEPA21Config):
        super().__init__()
        self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.cross_attn = VJEPA21PoolerCrossAttention(config)
        self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp = VJEPA21MLP(config, hidden_size=config.hidden_size)

    def forward(self, queries: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
        residual = queries
        hidden, _ = self.cross_attn(queries, self.layer_norm1(kv))
        hidden = residual + hidden

        residual = hidden
        hidden = residual + self.mlp(self.layer_norm2(hidden))
        return hidden


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))
        self.cross_attention_layer = VJEPA21PoolerCrossAttentionLayer(config)
        self.self_attention_layers = nn.ModuleList(
            [VJEPA21PoolerSelfAttentionLayer(config) for _ in range(config.num_pooler_layers)]
        )

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        for layer in self.self_attention_layers:
            hidden_state = layer(hidden_state)
        queries = self.query_tokens.expand(hidden_state.shape[0], -1, -1)
        return self.cross_attention_layer(queries, hidden_state).squeeze(1)


# ---------------------------------------------------------------------------
# PreTrainedModel base
# ---------------------------------------------------------------------------


class VJEPA21PreTrainedModel(PreTrainedModel):
    config_class = VJEPA21Config
    base_model_prefix = "vjepa21"
    main_input_name = "pixel_values_videos"
    supports_gradient_checkpointing = True
    _no_split_modules = ["VJEPA21Layer"]
    _supports_sdpa = True
    _supports_flash_attn = True
    _supports_flash_attn_2 = True

    @torch.no_grad()
    def _init_weights(self, module: nn.Module):
        std = self.config.initializer_range

        if isinstance(module, VJEPA21AttentivePooler):
            nn.init.trunc_normal_(module.query_tokens, std=std)
        elif isinstance(module, VJEPA21PredictorEmbeddings):
            if self.config.pred_zero_init_mask_tokens:
                for mt in module.mask_tokens:
                    nn.init.zeros_(mt)
            else:
                for mt in module.mask_tokens:
                    nn.init.trunc_normal_(mt, std=std)
        elif isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)):
            nn.init.trunc_normal_(module.weight, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.zeros_(module.bias)
            nn.init.ones_(module.weight)

        if isinstance(module, VJEPA21Embeddings):
            if module.img_mod_embed is not None:
                nn.init.normal_(module.img_mod_embed, std=1e-6)
                nn.init.normal_(module.video_mod_embed, std=1e-6)
        if isinstance(module, VJEPA21PredictorEmbeddings):
            if module.img_mod_embed is not None:
                nn.init.normal_(module.img_mod_embed, std=1e-6)
                nn.init.normal_(module.video_mod_embed, std=1e-6)


# ---------------------------------------------------------------------------
# Main Models
# ---------------------------------------------------------------------------


class VJEPA21Model(VJEPA21PreTrainedModel):
    """V-JEPA 2.1 model (encoder + predictor).

    Example, feature extraction:

    ```python
    model = VJEPA21Model.from_pretrained("apiantonio/vjepa2.1-vit-base-384")
    outputs = model(pixel_values_videos, skip_predictor=True)
    features = outputs.last_hidden_state
    ```

    Example, multi-level features (the `out_layers` recipe of the official
    frozen-evaluation probes):

    ```python
    levels = model(
        pixel_values_videos,
        skip_predictor=True,
        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):
        super().__init__(config)
        self.encoder = VJEPA21Encoder(config)
        self.predictor = VJEPA21Predictor(config)
        self.post_init()

    def get_input_embeddings(self):
        return self.encoder.embeddings.patch_embeddings

    def forward(
        self,
        pixel_values_videos: torch.Tensor,
        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,
    ) -> VJEPA21ModelOutput:
        """
        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 when `masks` is None.
            target_mask: List of `(B, K)` index tensors for target 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
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            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.
        needs_hierarchical_input = (
            not skip_predictor and len(self.config.predictor_hierarchical_layers) > 1
        )

        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,
        )
        seq_output = encoder_out.last_hidden_state

        predictor_output = None
        masked_hidden_state = None

        if not skip_predictor:
            predictor_input = (
                encoder_out.hierarchical_hidden_state
                if needs_hierarchical_input
                else seq_output
            )
            batch_size = seq_output.size(0)
            num_tokens = seq_output.size(1)
            device = seq_output.device

            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(
                predictor_input,
                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 = (
                predictor_input
                if context_is_masked
                else apply_masks(predictor_input, context_mask)
            )

        return VJEPA21ModelOutput(
            last_hidden_state=seq_output,
            hierarchical_hidden_state=(
                encoder_out.hierarchical_hidden_state if return_hierarchical else None
            ),
            multilevel_hidden_states=encoder_out.multilevel_hidden_states,
            masked_hidden_state=masked_hidden_state,
            predictor_output=predictor_output,
            hidden_states=encoder_out.hidden_states,
            attentions=encoder_out.attentions,
        )

    def _detect_mode(self, pixel_values_videos: torch.Tensor) -> str:
        x = normalize_video_layout(pixel_values_videos, self.config.in_chans)
        T = x.shape[2]
        if self.config.img_temporal_dim_size is not None and T == self.config.img_temporal_dim_size:
            return "img"
        return "video"

    def get_vision_features(self, pixel_values_videos: torch.Tensor) -> torch.Tensor:
        """Extract encoder features (convenience method for VLM integration)."""
        return self.forward(pixel_values_videos, skip_predictor=True).last_hidden_state


class VJEPA21ForVideoClassification(VJEPA21PreTrainedModel):
    """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)
        self.num_labels = config.num_labels
        self.vjepa21 = VJEPA21Model(config)
        self.pooler = VJEPA21AttentivePooler(config)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        self.post_init()

    def forward(
        self,
        pixel_values_videos: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        **kwargs,
    ) -> ImageClassifierOutput:
        r"""
        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,
            **forwarded,
        )
        pooled = self.pooler(outputs.last_hidden_state)
        logits = self.classifier(pooled)

        loss = None
        if labels is not None:
            loss = self.loss_function(pooled_logits=logits, labels=labels, config=self.config)

        return ImageClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = [
    "VJEPA21Model",
    "VJEPA21PreTrainedModel",
    "VJEPA21ForVideoClassification",
]