xwm.masking

What a JEPA predicts: 2-D blocks, 3-D tubes, temporal splits. Masks are batch-shared int32 index arrays with static shapes, so a training step compiles once. See Masking.

Mask samplers: the choice of what a predictive world model predicts.

All samplers share one signature -- sampler(key) -> MaskBatch -- and return static shapes, so a sampled mask can be fed straight into a jitted step.

Modules:

Name	Description
base	Mask representation and shared helpers.
block	Multi-block masking for images (I-JEPA).
random	Uniform random masking.
temporal	Temporal context/target splits for action-conditioned prediction.
tube	Tube masking for video (V-JEPA).

Classes:

Name	Description
MaskBatch	A context/target split of one token grid.
MultiBlockMask2d	I-JEPA style multi-block sampler over a 2-D patch grid.
RandomMask	Sample a uniformly random context set; predict everything else.
TemporalSplit	Context is the first n_context_frames; targets are the frames after.
TubeMask3d	V-JEPA style tube sampler over a (gt, gh, gw) token grid.

Functions:

Name	Description
block_shapes	All integer (h, w) with h * w == area, aspect in range, fitting bounds.
boolean_mask	Convert an index array to a boolean (n_tokens,) map.
complement	Indices of range(n_tokens) not present in idx.
gather	Select rows of an (N, D) token sequence: returns (len(idx), D).
subsample	Draw exactly size indices from idx without replacement.
random_split	Partition n_tokens into n_context visible and the rest masked.
frame_indices	Flat token indices belonging to temporal position t.
long_range_tubes	V-JEPA long-range preset: 2 tubes covering ~40% of the spatial grid each.
short_range_tubes	V-JEPA short-range preset: 8 tubes covering ~15% of the spatial grid each.

MaskBatch

Bases: NamedTuple

A context/target split of one token grid.

Attributes:

Name	Type	Description
context	Array	(K_ctx,) indices the encoder is allowed to see.
targets	Array	(M, K_tgt) indices to predict, one row per target block. Multiple blocks let one context encoding be reused for several prediction problems, which is where most of I-JEPA's efficiency comes from.

MultiBlockMask2d

MultiBlockMask2d(grid: tuple[int, int], *, n_targets: int = 4, target_scale: float = 0.15, aspect_range: tuple[float, float] = (0.75, 1.5), context_scale: float | None = None, max_tries: int = 100)

Bases: Module

I-JEPA style multi-block sampler over a 2-D patch grid.

Parameters:

Name	Type	Description	Default
grid	tuple[int, int]	(gh, gw) patch grid, e.g. PatchEmbed2d.grid.	required
n_targets	int	number of target blocks to predict (I-JEPA uses 4).	4
target_scale	float	fraction of the grid covered by each target block.	0.15
aspect_range	tuple[float, float]	allowed width/height ratios for target blocks.	(0.75, 1.5)
context_scale	float | None	fraction of all tokens kept as context. None derives a value the sampler can reliably satisfy given the expected overlap between target blocks.	None
max_tries	int	resampling budget when a draw leaves too little context.	100

Every target block holds exactly target_size tokens -- the closest usable size to round(target_scale * n_tokens). The aspect ratio is redrawn each call from the divisor pairs of that area, so block geometry varies while every returned array's shape stays static, which is what lets the training step be jitted once.

Source code in xwm/masking/block.py

def __init__(
    self,
    grid: tuple[int, int],
    *,
    n_targets: int = 4,
    target_scale: float = 0.15,
    aspect_range: tuple[float, float] = (0.75, 1.5),
    context_scale: float | None = None,
    max_tries: int = 100,
):
    gh, gw = grid
    n_tokens = gh * gw
    requested = max(1, int(round(target_scale * n_tokens)))
    target_size, shapes = snap_area(requested, aspect_range, (gh, gw))
    if context_scale is None:
        context_scale = expected_context_fraction(target_size / n_tokens, n_targets)
    context_size = max(1, int(round(context_scale * n_tokens)))
    if context_size >= n_tokens:
        raise ValueError("context_scale must leave at least one token masked")
    self.grid = (gh, gw)
    self.n_targets = n_targets
    self.target_size = target_size
    self.context_size = context_size
    self.shapes = tuple(shapes)
    self.max_tries = max_tries

RandomMask

RandomMask(n_tokens: int, mask_ratio: float = 0.75, n_targets: int = 1)

Bases: Module

Sample a uniformly random context set; predict everything else.

Parameters:

Name	Type	Description	Default
n_tokens	int	size of the token grid.	required
mask_ratio	float	fraction of tokens to predict.	0.75
n_targets	int	split the masked tokens into this many equal target blocks.	1

Source code in xwm/masking/random.py

def __init__(self, n_tokens: int, mask_ratio: float = 0.75, n_targets: int = 1):
    n_masked = int(round(mask_ratio * n_tokens))
    if not 0 < n_masked < n_tokens:
        raise ValueError(f"mask_ratio={mask_ratio} leaves nothing to predict or to see")
    if n_masked % n_targets:
        n_masked -= n_masked % n_targets  # keep target blocks equal-sized
    self.n_tokens = n_tokens
    self.mask_ratio = n_masked / n_tokens
    self.n_targets = n_targets

TemporalSplit

TemporalSplit(grid: tuple[int, int, int], *, n_context_frames: int = 1, horizon: int | None = None)

Bases: Module

Context is the first n_context_frames; targets are the frames after.

Parameters:

Name	Type	Description	Default
grid	tuple[int, int, int]	(gt, gh, gw) token grid.	required
n_context_frames	int	number of leading temporal positions kept visible.	1
horizon	int | None	how many future frames to predict; None uses all remaining.	None

Deterministic -- the key argument exists only so every sampler in xwm shares one call signature.

Source code in xwm/masking/temporal.py

def __init__(
    self,
    grid: tuple[int, int, int],
    *,
    n_context_frames: int = 1,
    horizon: int | None = None,
):
    gt = grid[0]
    max_horizon = gt - n_context_frames
    if max_horizon < 1:
        raise ValueError(f"grid has {gt} temporal positions, nothing left to predict")
    self.grid = tuple(grid)
    self.n_context_frames = n_context_frames
    self.horizon = max_horizon if horizon is None else min(horizon, max_horizon)

TubeMask3d

TubeMask3d(grid: tuple[int, int, int], *, n_targets: int = 8, spatial_scale: float = 0.15, temporal_extent: int | None = None, aspect_range: tuple[float, float] = (0.75, 1.5), context_scale: float | None = None, max_tries: int = 100)

Bases: Module

V-JEPA style tube sampler over a (gt, gh, gw) token grid.

Parameters:

Name	Type	Description	Default
grid	tuple[int, int, int]	token grid, e.g. PatchEmbed3d.grid.	required
n_targets	int	number of tubes to predict.	8
spatial_scale	float	fraction of the spatial grid each tube covers.	0.15
temporal_extent	int | None	tube length in temporal tokens; None spans the whole clip (the V-JEPA default).	None
aspect_range	tuple[float, float]	allowed width/height ratios for the tube cross-section.	(0.75, 1.5)
context_scale	float | None	fraction of all tokens kept as context. None derives a reliably satisfiable value from the expected tube overlap.	None
max_tries	int	resampling budget when a draw leaves too little context.	100

Attributes:

Name	Type	Description
target_size	int	Tokens per tube: cross-section area times temporal extent.

Source code in xwm/masking/tube.py

def __init__(
    self,
    grid: tuple[int, int, int],
    *,
    n_targets: int = 8,
    spatial_scale: float = 0.15,
    temporal_extent: int | None = None,
    aspect_range: tuple[float, float] = (0.75, 1.5),
    context_scale: float | None = None,
    max_tries: int = 100,
):
    gt, gh, gw = grid
    n_frames = gt if temporal_extent is None else min(temporal_extent, gt)
    requested = max(1, int(round(spatial_scale * gh * gw)))
    spatial_size, shapes = snap_area(requested, aspect_range, (gh, gw))
    n_tokens = gt * gh * gw
    if context_scale is None:
        # A tube spanning the full clip covers `spatial_size / (gh * gw)` of
        # every frame, so coverage is measured on the spatial grid alone.
        coverage = (spatial_size / (gh * gw)) * (n_frames / gt)
        context_scale = expected_context_fraction(coverage, n_targets)
    context_size = max(1, int(round(context_scale * n_tokens)))
    if context_size >= n_tokens:
        raise ValueError("context_scale must leave at least one token masked")
    self.grid = (gt, gh, gw)
    self.n_targets = n_targets
    self.temporal_extent = n_frames
    self.spatial_size = spatial_size
    self.context_size = context_size
    self.shapes = tuple(shapes)
    self.max_tries = max_tries

target_size

target_size: int

Tokens per tube: cross-section area times temporal extent.

block_shapes

block_shapes(area: int, aspect_range: tuple[float, float], bounds: tuple[int, int] | None = None) -> list[tuple[int, int]]

All integer (h, w) with h * w == area, aspect in range, fitting bounds.

Fixing the area while varying the aspect ratio is how xwm gets I-JEPA's varied block geometry without varying array shapes: every candidate yields exactly area tokens, so the sampler's output shape is static.

Source code in xwm/masking/base.py

def block_shapes(
    area: int,
    aspect_range: tuple[float, float],
    bounds: tuple[int, int] | None = None,
) -> list[tuple[int, int]]:
    """All integer ``(h, w)`` with ``h * w == area``, aspect in range, fitting ``bounds``.

    Fixing the *area* while varying the aspect ratio is how xwm gets I-JEPA's
    varied block geometry without varying array shapes: every candidate yields
    exactly ``area`` tokens, so the sampler's output shape is static.
    """
    lo, hi = aspect_range
    divisors = (h for h in range(1, area + 1) if area % h == 0)
    # Aspect ratio of an h x (area / h) block is (area / h) / h.
    out = [(h, area // h) for h in divisors if lo <= (area / h) / h <= hi]
    if bounds is not None:
        out = [(h, w) for h, w in out if h <= bounds[0] and w <= bounds[1]]
    return out

boolean_mask

boolean_mask(idx: Array, n_tokens: int) -> Array

Convert an index array to a boolean (n_tokens,) map.

Source code in xwm/masking/base.py

def boolean_mask(idx: Array, n_tokens: int) -> Array:
    """Convert an index array to a boolean ``(n_tokens,)`` map."""
    return jnp.zeros((n_tokens,), dtype=bool).at[idx].set(True)

complement

complement(idx: ndarray, n_tokens: int) -> ndarray

Indices of range(n_tokens) not present in idx.

Source code in xwm/masking/base.py

def complement(idx: np.ndarray, n_tokens: int) -> np.ndarray:
    """Indices of ``range(n_tokens)`` not present in ``idx``."""
    keep = np.ones((n_tokens,), dtype=bool)
    keep[idx] = False
    return np.flatnonzero(keep)

gather

gather(tokens: Array, idx: Array) -> Array

Select rows of an (N, D) token sequence: returns (len(idx), D).

Source code in xwm/masking/base.py

def gather(tokens: Array, idx: Array) -> Array:
    """Select rows of an ``(N, D)`` token sequence: returns ``(len(idx), D)``."""
    return tokens[idx]

subsample

subsample(rng: Generator, idx: ndarray, size: int) -> ndarray

Draw exactly size indices from idx without replacement.

Raises if idx is too small -- samplers are constructed so this cannot happen, and a loud failure beats a silently reshaped batch.

Source code in xwm/masking/base.py

def subsample(rng: np.random.Generator, idx: np.ndarray, size: int) -> np.ndarray:
    """Draw exactly ``size`` indices from ``idx`` without replacement.

    Raises if ``idx`` is too small -- samplers are constructed so this cannot
    happen, and a loud failure beats a silently reshaped batch.
    """
    if idx.size < size:
        raise ValueError(f"cannot draw {size} indices from a pool of {idx.size}")
    return np.sort(rng.choice(idx, size=size, replace=False))

random_split

random_split(key: PRNGKey, n_tokens: int, n_context: int) -> tuple[Array, Array]

Partition n_tokens into n_context visible and the rest masked.

Uses the argsort-of-noise trick, so both outputs have static shapes and the whole thing is jittable and differentiable-through-free.

Source code in xwm/masking/random.py

def random_split(key: PRNGKey, n_tokens: int, n_context: int) -> tuple[Array, Array]:
    """Partition ``n_tokens`` into ``n_context`` visible and the rest masked.

    Uses the argsort-of-noise trick, so both outputs have static shapes and the
    whole thing is jittable and differentiable-through-free.
    """
    order = jnp.argsort(jr.uniform(key, (n_tokens,)))
    return jnp.sort(order[:n_context]), jnp.sort(order[n_context:])

frame_indices

frame_indices(grid: tuple[int, int, int], t: int) -> ndarray

Flat token indices belonging to temporal position t.

Source code in xwm/masking/temporal.py

def frame_indices(grid: tuple[int, int, int], t: int) -> np.ndarray:
    """Flat token indices belonging to temporal position ``t``."""
    _, gh, gw = grid
    return np.arange(t * gh * gw, (t + 1) * gh * gw)

long_range_tubes

long_range_tubes(grid: tuple[int, int, int], **kwargs) -> TubeMask3d

V-JEPA long-range preset: 2 tubes covering ~40% of the spatial grid each.

Source code in xwm/masking/tube.py

def long_range_tubes(grid: tuple[int, int, int], **kwargs) -> TubeMask3d:
    """V-JEPA long-range preset: 2 tubes covering ~40% of the spatial grid each."""
    kwargs.setdefault("n_targets", 2)
    kwargs.setdefault("spatial_scale", 0.4)
    return TubeMask3d(grid, **kwargs)

short_range_tubes

short_range_tubes(grid: tuple[int, int, int], **kwargs) -> TubeMask3d

V-JEPA short-range preset: 8 tubes covering ~15% of the spatial grid each.

Source code in xwm/masking/tube.py

def short_range_tubes(grid: tuple[int, int, int], **kwargs) -> TubeMask3d:
    """V-JEPA short-range preset: 8 tubes covering ~15% of the spatial grid each."""
    kwargs.setdefault("n_targets", 8)
    kwargs.setdefault("spatial_scale", 0.15)
    return TubeMask3d(grid, **kwargs)