torch_world.nn¶
Encoders¶
CLSEncoder(
img_size: int,
patch_size: int,
embed_dim: int = 192,
depth: int = 6,
num_heads: int = 3,
resize: bool = False,
**kwargs,
)
Bases: Module
A ViT whose CLS token embeds each image: (N, C, H, W) to (N, embed_dim).
With resize, inputs of any size are resized to img_size first, as for the local crops of multi-crop
training. Other keyword arguments go to VisionTransformer.
Bases: Module
Raw pixel patches as tokens: (N, C, H, W) to (N, H/p * W/p, C * p * p), with no parameters.
A fast, self-contained stand-in for a frozen pre-trained encoder such as DINOv2 or V-JEPA. Because the features
are pixels, unpatchify turns predicted tokens back into images.
Tokens (..., P, C * p * p) on a grid of patches back to images (..., C, H, W).
Bases: Module
Wraps a DINOv2 model to return its normalized patch tokens (N, P, D).
Inputs are float images in [0, 1] or uint8 images. size resizes them first, and normalize applies the
ImageNet mean and standard deviation that DINOv2 expects. PatchTokens.from_hub loads a released model.
Oquab et al., "DINOv2: Learning Robust Visual Features without Supervision", TMLR 2024. Paper: https://arxiv.org/abs/2304.07193 Code: https://github.com/facebookresearch/dinov2
classmethod
¶
from_hub(
name: str = "dinov2_vits14", size: int | None = 224
) -> PatchTokens
A frozen, normalized released DINOv2 (dinov2_vits14, _vitb14, _vitl14 or _vitg14).
Bases: Module
Token embedding, a bidirectional transformer and mean pooling over non-padding tokens: (B, dim).
A small stand-in for a pretrained text encoder such as the EmbeddingGemma Y-encoder of VL-JEPA.
Whitespace tokenizer over a fixed vocabulary, for small closed-vocabulary text such as templated QA.
Id 0 is padding. tokenizer(texts) returns (ids, mask) of shape (B, L), padded to the longest text.
With a pretrained text model, use its own tokenizer instead.
Transformers¶
VisionTransformer(
img_size: int | tuple[int, int] = 224,
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768,
depth: int = 12,
num_heads: int = 12,
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
drop_path: float = 0.0,
num_frames: int = 1,
tubelet_size: int = 2,
cls_token: bool = False,
use_rope: bool = False,
init_std: float = 0.02,
)
Bases: Module
ViT for images or videos (num_frames > 1 uses tubelet_size frames per token).
Positions are fixed sin-cos embeddings, or 2D/3D RoPE when use_rope. forward takes optional
index masks and then encodes only the kept tokens, as JEPA context encoders do.
classmethod
¶
from_size(size: str, **kwargs) -> VisionTransformer
Build a standard ViT: tiny, small, base, large, huge or giant.
PatchEmbed(
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768,
tubelet_size: int | None = None,
)
Bases: Module
Images (B, C, H, W) or videos (B, C, T, H, W) to patch tokens (B, N, D).
MaskedPredictor(
grid: tuple[int, ...],
embed_dim: int,
predictor_dim: int = 384,
depth: int = 6,
num_heads: int = 12,
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
drop_path: float = 0.0,
num_mask_tokens: int = 1,
zero_init_mask_tokens: bool = False,
use_rope: bool = False,
init_std: float = 0.02,
)
Bases: Module
JEPA predictor: from context tokens, predict target tokens at masked positions of a patch grid.
Each target position is a learnable mask token plus its position. num_mask_tokens holds one token
per masking strategy, as in V-JEPA.
x is (len(masks_x) * B, Nx, D); returns (len(masks_y) * len(masks_x) * B, Ny, D).
Transformer(
dim: int,
depth: int,
num_heads: int,
drop_path: float = 0.0,
conditional: bool = False,
**kwargs,
)
Bases: Module
Stack of Block (or AdaLNBlock when conditional) followed by a final layer norm.
Block(
dim: int,
num_heads: int,
mlp_dim: int | None = None,
head_dim: int | None = None,
qkv_bias: bool = True,
dropout: float = 0.0,
drop_path: float = 0.0,
use_rope: bool = False,
act: Callable[[], Module] = GELU,
)
Bases: Module
Pre-norm transformer block.
Bases: Block
Transformer block conditioned by adaptive layer norm with zero-initialized gates.
Peebles & Xie, "Scalable Diffusion Models with Transformers", ICCV 2023. Paper: https://arxiv.org/abs/2212.09748
Attention(
dim: int,
num_heads: int = 8,
head_dim: int | None = None,
qkv_bias: bool = True,
dropout: float = 0.0,
use_rope: bool = False,
)
Bases: Module
Multi-head self-attention on F.scaled_dot_product_attention with optional rotary positions.
forward(
x: Tensor,
attn_mask: Tensor | None = None,
pos: Tensor | None = None,
is_causal: bool = False,
) -> Tensor
attn_mask is boolean (True attends) or additive; pos holds RoPE coordinates (N, A) or (B, N, A).
MLP(
dims: list[int],
norm: Callable[[int], Module] | None = None,
act: Callable[[], Module] = GELU,
dropout: float = 0.0,
)
Bases: Sequential
Linear layers through dims with norm, act and dropout between them.
Positions¶
Fixed 2D sin-cos embedding of shape (H * W, dim), as in V-JEPA.
Fixed 3D sin-cos embedding of shape (T * H * W, dim) with uniform power per axis, as in V-JEPA.
Integer coordinates of every cell of grid in row-major order, shape (prod(grid), len(grid)).
Rotary embedding over pos.size(-1) axes, each rotating an equal even slice of the last dim of x.
x is (..., N, D) and pos is (N, A) or broadcastable to (..., N, A).
Dimensions left after A equal slices are not rotated.
Initialization¶
Truncated-normal linear weights, zero biases and unit norms, as in the JEPA reference code.
Divide residual output projections of layer i by sqrt(2 i).