mstar.model.vjepa2.components.rope_utils

mstar.model.vjepa2.components.rope_utils#

3D rotary position embedding helpers for V-JEPA 2.

Ported verbatim from HuggingFace transformers/models/vjepa2/modeling_vjepa2.py and upstream vjepa2/src/models/utils/modules.py. Both copies share the same frequency schedule and the same “duplicate then rotate adjacent pair” construction noted in the upstream comment as a pretraining-compatibility quirk. Do not “fix” the .repeat(…, 2) to .repeat_interleave(2) — it would break the pretrained weights.

Functions

rotate_queries_or_keys(x, pos)

Apply rotary embeddings along the last D dims of x.

rotate_queries_or_keys_BNHD(x, pos)

mstar.model.vjepa2.components.rope_utils.rotate_queries_or_keys(x, pos)[source]#

Apply rotary embeddings along the last D dims of x.

Parameters:

  • x (Tensor) – [B, num_heads, N, D] (or broadcastable) — queries or keys.

  • pos (Tensor) – positions of shape [N] or [B, num_heads, N].

Returns:

Rotated tensor, same shape as x.

Return type:

Tensor

mstar.model.vjepa2.components.rope_utils.rotate_queries_or_keys_BNHD(x, pos)[source]#
Parameters:
Return type:

Tensor