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

apply_rope_BNHD(x, emb_sin, emb_cos)

rope_tables_BNHD(pos, dim, dtype, device)

(sin, cos) for apply_rope_BNHD over the last dim dims.

rotate_queries_or_keys(x, pos)

Apply rotary embeddings along the last D dims of x.

Classes

RopeTables(spatial, time)

One cached step's (sin, cos) (ACRoPEAttention.rope_tables): built once from the step's positions and shared by every block, q and k.
class mstar.model.vjepa2.components.rope_utils.RopeTables(spatial, time)[source]#

Bases: NamedTuple

One cached step’s (sin, cos) (ACRoPEAttention.rope_tables): built once from the step’s positions and shared by every block, q and k.

Parameters:
spatial: tuple[Tensor, Tensor]#

Alias for field number 0

time: tuple[Tensor, Tensor] | None#

Alias for field number 1

mstar.model.vjepa2.components.rope_utils.apply_rope_BNHD(x, emb_sin, emb_cos)[source]#
Parameters:
Return type:

Tensor

mstar.model.vjepa2.components.rope_utils.rope_tables_BNHD(pos, dim, dtype, device)[source]#

(sin, cos) for apply_rope_BNHD over the last dim dims.

Parameters:
Return type:

tuple[Tensor, Tensor]

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