torch_world.masks¶

MultiBlockMask ¶

MultiBlockMask(
    grid: tuple[int, int],
    enc_scale: tuple[float, float] = (0.85, 1.0),
    pred_scale: tuple[float, float] = (0.15, 0.2),
    aspect_ratio: tuple[float, float] = (0.75, 1.5),
    num_enc: int = 1,
    num_pred: int = 4,
    min_keep: int = 10,
    allow_overlap: bool = False,
)

I-JEPA multi-block masking on a 2D patch grid.

Samples num_pred target blocks and num_enc context blocks per image. Context blocks avoid the targets unless allow_overlap. Block sizes are shared across the batch, and every mask is cropped to the shortest one so masks stack into (B, K) index tensors.

MultiBlock3DMask ¶

MultiBlock3DMask(
    grid: tuple[int, int, int],
    spatial_scale: tuple[float, float] = (0.15, 0.15),
    temporal_scale: tuple[float, float] = (1.0, 1.0),
    aspect_ratio: tuple[float, float] = (0.75, 1.5),
    num_blocks: int = 8,
    max_temporal_keep: float = 1.0,
)

V-JEPA spatio-temporal multi-block masking on a (T, H, W) token grid.

The union of num_blocks tubes is the prediction target; the remaining tokens are the context.

apply_masks ¶

apply_masks(x: Tensor, masks: list[Tensor]) -> Tensor

Gather the tokens indexed by each (B, K) mask from x of shape (B, N, ...); stack masks on the batch dim.

frame_causal_mask ¶

frame_causal_mask(
    num_frames: int,
    tokens_per_frame: int,
    window: int | None = None,
    device=None,
) -> Tensor

Boolean (N, N) mask where a token attends every token of its own frame and of the window frames before it.

repeat_interleave_batch ¶

repeat_interleave_batch(
    x: Tensor, batch_size: int, repeat: int
) -> Tensor

Repeat each consecutive chunk of batch_size rows repeat times.