mstar.model.vjepa2.components.predictor

V-JEPA 2 masked latent predictor.

Ports VJEPA2PredictorEmbeddings and VJEPA2Predictor from HuggingFace transformers/models/vjepa2/modeling_vjepa2.py.

Takes encoder hidden states + context_mask (positions the predictor sees) + target_mask (positions to predict) and emits predicted embeddings at the target positions. NOT autoregressive — a single forward predicts every target token in parallel.

Weight layout (prefix predictor.):

predictor.embeddings.predictor_embeddings.{weight,bias} predictor.embeddings.mask_tokens predictor.layer.{N}.* predictor.layernorm.{weight,bias} predictor.proj.{weight,bias}

Classes

VJEPA2Predictor(config)
VJEPA2PredictorEmbeddings(config)	Project encoder hidden states into predictor space and concatenate learned mask tokens at the target positions.

class mstar.model.vjepa2.components.predictor.VJEPA2Predictor(config)

Bases: Module

Parameters:

config (VJepa2Config)

forward(encoder_hidden_states, context_mask, target_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• encoder_hidden_states (Tensor)

• context_mask (list[Tensor])

• target_mask (list[Tensor])

Return type:

Tensor

class mstar.model.vjepa2.components.predictor.VJEPA2PredictorEmbeddings(config)

Bases: Module

Project encoder hidden states into predictor space and concatenate learned mask tokens at the target positions.

Parameters:

config (VJepa2Config)

forward(hidden_states, context_mask, target_mask, mask_index=1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• hidden_states (Tensor)

• context_mask (list[Tensor])

• target_mask (list[Tensor])

• mask_index (int)

Return type:

tuple[Tensor, Tensor]