xwm.nn

Layers: attention, transformer blocks, RoPE, patch and tubelet embeddings, normalisations. Unbatched, vmaped by the caller.

Generic neural building blocks: layers that know nothing about world models.

These are the pieces you compose into encoders and predictors. Every module is written for a single unbatched sample; vmap for batches.

Modules:

Name	Description
attention	Self- and cross-attention over unbatched token sequences (N, D).
drop	Stochastic-depth (drop-path) regularisation.
embed	Positional information: fixed sin-cos tables, learned tables, axial RoPE.
mlp	Feed-forward blocks.
norm	Normalisation layers.
patch	Patch and tubelet embeddings -- the entry point from pixels to tokens.
transformer	Pre-norm transformer blocks and stacks.

Classes:

Name	Description
Attention	Multi-head attention with optional QK-norm and axial RoPE.
CrossAttention	Attention from a query sequence into a separate key/value sequence.
DropPath	Drop a residual branch entirely, with probability p.
AxialRoPE	Axial rotary embeddings for 1-D, 2-D (image) or 3-D (video) grids.
LearnedPosEmbed	A trainable (N, D) position table.
SinCosPosEmbed	A frozen factorised sin-cos table over a grid.
Mlp	Two-layer position-wise MLP, applied to each token independently.
SwiGLU	Gated feed-forward (SwiGLU), a common drop-in for Mlp.
LayerNorm	Layer normalisation over the last axis, with optional affine params.
RMSNorm	Root-mean-square normalisation (no mean subtraction).
SimNorm	Simplicial normalisation: softmax over fixed-size groups of channels.
PatchEmbed2d	Split an image into non-overlapping patches and linearly embed them.
PatchEmbed3d	Tubelet embedding: patches spanning tubelet_size frames.
AttentivePooler	Pool a token sequence into n_queries vectors via cross-attention.
Block	Pre-norm self-attention block: x + attn(ln(x)), x + mlp(ln(x)).
CrossBlock	Pre-norm cross-attention block, then an MLP.
LayerScale	Per-channel learnable residual gain, initialised near zero.
Transformer	A stack of Block s with a final norm.

Functions:

Name	Description
attend	Scaled dot-product attention on (heads, N, head_dim) inputs.
block_causal_attention_mask	Causal across blocks, fully connected within a block.
causal_attention_mask	Boolean (n, n) mask allowing each position to see itself and the past.
linear_droppath_schedule	Linearly increasing drop-path rates across depth (the ViT default).
grid_coords	(prod(grid), len(grid)) integer coordinates, row-major flattened.
sincos_pos_embed	Fixed sin-cos table for an n-dimensional grid, flattened row-major.
l2_normalize	Project onto the unit sphere along axis.
unpatchify_2d	Fold (N, C * p * p) patch values back into a (C, H, W) image.
mean_pool	Average over the token axis of an (N, D) sequence.

Attention

Attention(dim: int, num_heads: int, *, key: PRNGKey | None = None, qkv_bias: bool = True, qk_norm: bool = False, dropout: float = 0.0, rope: AxialRoPE | None = None)

Bases: Module

Multi-head attention with optional QK-norm and axial RoPE.

Set rope to an AxialRoPE to use relative positions; leave it None and add an absolute table to the inputs instead (what I-JEPA and V-JEPA do).

Source code in xwm/nn/attention.py

def __init__(
    self,
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    qkv_bias: bool = True,
    qk_norm: bool = False,
    dropout: float = 0.0,
    rope: AxialRoPE | None = None,
):
    key = resolve_key(key)
    if dim % num_heads != 0:
        raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}")
    k1, k2 = jr.split(key)
    head_dim = dim // num_heads
    self.num_heads = num_heads
    self.qkv = eqx.nn.Linear(dim, 3 * dim, use_bias=qkv_bias, key=k1)
    self.proj = eqx.nn.Linear(dim, dim, key=k2)
    self.q_norm = LayerNorm(head_dim) if qk_norm else None
    self.k_norm = LayerNorm(head_dim) if qk_norm else None
    self.rope = rope
    self.drop = eqx.nn.Dropout(dropout)

CrossAttention

CrossAttention(dim: int, num_heads: int, *, key: PRNGKey | None = None, kv_dim: int | None = None, qkv_bias: bool = True, qk_norm: bool = False, dropout: float = 0.0)

Bases: Module

Attention from a query sequence into a separate key/value sequence.

Used by the attentive probe (a learned query pools an encoder's tokens) and by predictors that keep context tokens read-only.

Source code in xwm/nn/attention.py

def __init__(
    self,
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    kv_dim: int | None = None,
    qkv_bias: bool = True,
    qk_norm: bool = False,
    dropout: float = 0.0,
):
    key = resolve_key(key)
    if dim % num_heads != 0:
        raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}")
    kv_dim = kv_dim or dim
    k1, k2, k3 = jr.split(key, 3)
    head_dim = dim // num_heads
    self.num_heads = num_heads
    self.to_q = eqx.nn.Linear(dim, dim, use_bias=qkv_bias, key=k1)
    self.to_kv = eqx.nn.Linear(kv_dim, 2 * dim, use_bias=qkv_bias, key=k2)
    self.proj = eqx.nn.Linear(dim, dim, key=k3)
    self.q_norm = LayerNorm(head_dim) if qk_norm else None
    self.k_norm = LayerNorm(head_dim) if qk_norm else None
    self.drop = eqx.nn.Dropout(dropout)

DropPath

DropPath(p: float = 0.0, *, inference: bool = False)

Bases: Module

Drop a residual branch entirely, with probability p.

Applied per-sample; since xwm modules are unbatched this is a single Bernoulli draw that zeroes (or rescales) the whole branch.

Source code in xwm/nn/drop.py

def __init__(self, p: float = 0.0, *, inference: bool = False):
    self.p = p
    self.inference = inference

AxialRoPE

AxialRoPE(grid: Grid, head_dim: int, *, base: float = 100.0)

Bases: Module

Axial rotary embeddings for 1-D, 2-D (image) or 3-D (video) grids.

The head dimension is split across axes by axis_dims and each chunk is rotated by its axis' coordinate. Unlike additive tables, RoPE acts inside attention and encodes relative position, which extrapolates better to resolutions and clip lengths unseen during training.

Source code in xwm/nn/embed.py

def __init__(self, grid: Grid, head_dim: int, *, base: float = 100.0):
    self.grid = tuple(grid)
    self.head_dim = head_dim
    self.base = base
    _rope_tables(self.grid, head_dim, base)  # fail fast

LearnedPosEmbed

LearnedPosEmbed(n_positions: int, dim: int, *, key: PRNGKey | None = None, scale: float = 0.02)

Bases: Module

A trainable (N, D) position table.

Source code in xwm/nn/embed.py

def __init__(
    self,
    n_positions: int,
    dim: int,
    *,
    key: PRNGKey | None = None,
    scale: float = 0.02,
):
    self.table = scale * jr.normal(resolve_key(key), (n_positions, dim))

SinCosPosEmbed

SinCosPosEmbed(grid: Grid, dim: int)

Bases: Module

A frozen factorised sin-cos table over a grid.

Source code in xwm/nn/embed.py

def __init__(self, grid: Grid, dim: int):
    self.grid = tuple(grid)
    self.dim = dim
    sincos_pos_embed(self.grid, dim)  # fail fast on bad dims

Mlp

Mlp(dim: int, hidden_dim: int | None = None, out_dim: int | None = None, *, key: PRNGKey | None = None, act: Callable[[Array], Array] = gelu, dropout: float = 0.0)

Bases: Module

Two-layer position-wise MLP, applied to each token independently.

Source code in xwm/nn/mlp.py

def __init__(
    self,
    dim: int,
    hidden_dim: int | None = None,
    out_dim: int | None = None,
    *,
    key: PRNGKey | None = None,
    act: Callable[[Array], Array] = jax.nn.gelu,
    dropout: float = 0.0,
):
    key = resolve_key(key)
    hidden_dim = hidden_dim or 4 * dim
    out_dim = out_dim or dim
    k1, k2 = jr.split(key)
    self.fc1 = eqx.nn.Linear(dim, hidden_dim, key=k1)
    self.fc2 = eqx.nn.Linear(hidden_dim, out_dim, key=k2)
    self.act = act
    self.drop = eqx.nn.Dropout(dropout)

SwiGLU

SwiGLU(dim: int, hidden_dim: int | None = None, out_dim: int | None = None, *, key: PRNGKey | None = None, dropout: float = 0.0)

Bases: Module

Gated feed-forward (SwiGLU), a common drop-in for Mlp.

Source code in xwm/nn/mlp.py

def __init__(
    self,
    dim: int,
    hidden_dim: int | None = None,
    out_dim: int | None = None,
    *,
    key: PRNGKey | None = None,
    dropout: float = 0.0,
):
    # 2/3 keeps the parameter count comparable to a 4x dense MLP.
    key = resolve_key(key)
    hidden_dim = hidden_dim or int(8 * dim / 3)
    out_dim = out_dim or dim
    k1, k2, k3 = jr.split(key, 3)
    self.w_in = eqx.nn.Linear(dim, hidden_dim, key=k1)
    self.w_gate = eqx.nn.Linear(dim, hidden_dim, key=k2)
    self.w_out = eqx.nn.Linear(hidden_dim, out_dim, key=k3)
    self.drop = eqx.nn.Dropout(dropout)

LayerNorm

LayerNorm(dim: int, *, eps: float = 1e-06, affine: bool = True)

Bases: Module

Layer normalisation over the last axis, with optional affine params.

Source code in xwm/nn/norm.py

def __init__(self, dim: int, *, eps: float = 1e-6, affine: bool = True):
    self.weight = jnp.ones((dim,)) if affine else None
    self.bias = jnp.zeros((dim,)) if affine else None
    self.eps = eps

RMSNorm

RMSNorm(dim: int, *, eps: float = 1e-06, affine: bool = True)

Bases: Module

Root-mean-square normalisation (no mean subtraction).

Source code in xwm/nn/norm.py

def __init__(self, dim: int, *, eps: float = 1e-6, affine: bool = True):
    self.weight = jnp.ones((dim,)) if affine else None
    self.eps = eps

SimNorm

SimNorm(group_size: int = 8)

Bases: Module

Simplicial normalisation: softmax over fixed-size groups of channels.

TD-MPC2's latent normalisation, and the detail its stability rests on. The latent is split into groups of group_size channels and each group is softmaxed, so the latent becomes a concatenation of points on probability simplices.

Why that helps where LayerNorm does not: the representation is bounded (every entry in [0, 1], every group summing to one), so a long recurrent unroll cannot drift to infinity -- but it also cannot collapse to a single point, because each group must keep its mass distributed to stay off the simplex corners. Bounded without being degenerate is exactly what a model that gets unrolled inside a planner needs.

Source code in xwm/nn/norm.py

def __init__(self, group_size: int = 8):
    if group_size < 2:
        raise ValueError(f"group_size must be at least 2, got {group_size}")
    self.group_size = group_size

PatchEmbed2d

PatchEmbed2d(img_size: int | tuple[int, int], patch_size: int, in_channels: int, embed_dim: int, *, key: PRNGKey | None = None, bias: bool = True)

Bases: Module

Split an image into non-overlapping patches and linearly embed them.

Implemented as a reshape plus a matmul rather than a strided convolution: identical arithmetic, but it makes the token ordering explicit (row-major over (H // p, W // p)), which the masking code depends on.

Source code in xwm/nn/patch.py

def __init__(
    self,
    img_size: int | tuple[int, int],
    patch_size: int,
    in_channels: int,
    embed_dim: int,
    *,
    key: PRNGKey | None = None,
    bias: bool = True,
):
    key = resolve_key(key)
    h, w = (img_size, img_size) if isinstance(img_size, int) else img_size
    if h % patch_size or w % patch_size:
        raise ValueError(f"img_size {(h, w)} not divisible by patch_size {patch_size}")
    fan_in = in_channels * patch_size * patch_size
    k1, k2 = jr.split(key)
    self.weight = jr.normal(k1, (fan_in, embed_dim)) * fan_in**-0.5
    self.bias = jnp.zeros((embed_dim,)) if bias else None
    self.img_size = (h, w)
    self.patch_size = patch_size
    self.in_channels = in_channels
    self.embed_dim = embed_dim

PatchEmbed3d

PatchEmbed3d(img_size: int | tuple[int, int], patch_size: int, num_frames: int, tubelet_size: int, in_channels: int, embed_dim: int, *, key: PRNGKey | None = None, bias: bool = True)

Bases: Module

Tubelet embedding: patches spanning tubelet_size frames.

Video JEPAs tokenise space and time jointly, so a clip of T frames at tubelet_size = 2 yields T // 2 temporal positions. Token order is row-major over (T // ts, H // p, W // p).

Source code in xwm/nn/patch.py

def __init__(
    self,
    img_size: int | tuple[int, int],
    patch_size: int,
    num_frames: int,
    tubelet_size: int,
    in_channels: int,
    embed_dim: int,
    *,
    key: PRNGKey | None = None,
    bias: bool = True,
):
    key = resolve_key(key)
    h, w = (img_size, img_size) if isinstance(img_size, int) else img_size
    if h % patch_size or w % patch_size:
        raise ValueError(f"img_size {(h, w)} not divisible by patch_size {patch_size}")
    if num_frames % tubelet_size:
        raise ValueError(f"num_frames {num_frames} not divisible by tubelet {tubelet_size}")
    fan_in = in_channels * tubelet_size * patch_size * patch_size
    k1, _ = jr.split(key)
    self.weight = jr.normal(k1, (fan_in, embed_dim)) * fan_in**-0.5
    self.bias = jnp.zeros((embed_dim,)) if bias else None
    self.img_size = (h, w)
    self.num_frames = num_frames
    self.patch_size = patch_size
    self.tubelet_size = tubelet_size
    self.in_channels = in_channels
    self.embed_dim = embed_dim

AttentivePooler

AttentivePooler(dim: int, num_heads: int, *, key: PRNGKey | None = None, kv_dim: int | None = None, n_queries: int = 1, depth: int = 1, mlp_ratio: float = 4.0)

Bases: Module

Pool a token sequence into n_queries vectors via cross-attention.

The standard read-out head for frozen JEPA encoders: a learned query attends over the tokens, which recovers substantially more of the representation than mean-pooling for the same probe budget.

Source code in xwm/nn/transformer.py

def __init__(
    self,
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    kv_dim: int | None = None,
    n_queries: int = 1,
    depth: int = 1,
    mlp_ratio: float = 4.0,
):
    key = resolve_key(key)
    kq, *kb = jr.split(key, depth + 1)
    self.query = 0.02 * jr.normal(kq, (n_queries, dim))
    self.blocks = [
        CrossBlock(dim, num_heads, key=kb[i], kv_dim=kv_dim, mlp_ratio=mlp_ratio)
        for i in range(depth)
    ]
    self.norm = LayerNorm(dim)

Block

Block(dim: int, num_heads: int, *, key: PRNGKey | None = None, mlp_ratio: float = 4.0, qkv_bias: bool = True, qk_norm: bool = False, dropout: float = 0.0, drop_path: float = 0.0, layer_scale: float | None = None, rope: AxialRoPE | None = None, act: Callable[[Array], Array] = gelu)

Bases: Module

Pre-norm self-attention block: x + attn(ln(x)), x + mlp(ln(x)).

Source code in xwm/nn/transformer.py

def __init__(
    self,
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_norm: bool = False,
    dropout: float = 0.0,
    drop_path: float = 0.0,
    layer_scale: float | None = None,
    rope: AxialRoPE | None = None,
    act: Callable[[Array], Array] = jax.nn.gelu,
):
    key = resolve_key(key)
    ka, km = jr.split(key)
    self.norm1 = LayerNorm(dim)
    self.attn = Attention(
        dim, num_heads, key=ka, qkv_bias=qkv_bias, qk_norm=qk_norm, dropout=dropout, rope=rope
    )
    self.norm2 = LayerNorm(dim)
    self.mlp = Mlp(dim, int(dim * mlp_ratio), key=km, act=act, dropout=dropout)
    self.ls1 = LayerScale(dim, layer_scale) if layer_scale else None
    self.ls2 = LayerScale(dim, layer_scale) if layer_scale else None
    self.dp1 = DropPath(drop_path)
    self.dp2 = DropPath(drop_path)

CrossBlock

CrossBlock(dim: int, num_heads: int, *, key: PRNGKey | None = None, kv_dim: int | None = None, mlp_ratio: float = 4.0, dropout: float = 0.0)

Bases: Module

Pre-norm cross-attention block, then an MLP.

Source code in xwm/nn/transformer.py

def __init__(
    self,
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    kv_dim: int | None = None,
    mlp_ratio: float = 4.0,
    dropout: float = 0.0,
):
    key = resolve_key(key)
    ka, km = jr.split(key)
    self.norm_q = LayerNorm(dim)
    self.norm_kv = LayerNorm(kv_dim or dim)
    self.attn = CrossAttention(dim, num_heads, key=ka, kv_dim=kv_dim, dropout=dropout)
    self.norm2 = LayerNorm(dim)
    self.mlp = Mlp(dim, int(dim * mlp_ratio), key=km, dropout=dropout)

LayerScale

LayerScale(dim: int, init: float = 0.0001)

Bases: Module

Per-channel learnable residual gain, initialised near zero.

Stabilises deep ViTs by starting each residual branch as a near-no-op.

Source code in xwm/nn/transformer.py

def __init__(self, dim: int, init: float = 1e-4):
    self.gamma = jnp.full((dim,), init)

Transformer

Transformer(dim: int, depth: int, num_heads: int, *, key: PRNGKey | None = None, mlp_ratio: float = 4.0, qkv_bias: bool = True, qk_norm: bool = False, dropout: float = 0.0, drop_path: float = 0.0, layer_scale: float | None = None, rope: AxialRoPE | None = None, final_norm: bool = True, remat: bool = False)

Bases: Module

A stack of Block s with a final norm.

Parameters:

Name	Type	Description	Default
depth	int	number of blocks.	required
drop_path	float	maximum stochastic-depth rate; rates increase linearly with depth, as in the ViT/DeiT recipe.	0.0
remat	bool	wrap each block in jax.checkpoint, trading recomputation for activation memory. Worth enabling for long video sequences.	False

Source code in xwm/nn/transformer.py

def __init__(
    self,
    dim: int,
    depth: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_norm: bool = False,
    dropout: float = 0.0,
    drop_path: float = 0.0,
    layer_scale: float | None = None,
    rope: AxialRoPE | None = None,
    final_norm: bool = True,
    remat: bool = False,
):
    key = resolve_key(key)
    rates = linear_droppath_schedule(depth, drop_path)
    keys = jr.split(key, depth)
    self.blocks = [
        Block(
            dim,
            num_heads,
            key=keys[i],
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            dropout=dropout,
            drop_path=rates[i],
            layer_scale=layer_scale,
            rope=rope,
        )
        for i in range(depth)
    ]
    self.norm = LayerNorm(dim) if final_norm else None
    self.remat = remat

attend

attend(q: Array, k: Array, v: Array, mask: Array | None = None) -> Array

Scaled dot-product attention on (heads, N, head_dim) inputs.

mask is a boolean (N_q, N_kv) array where True means attend. Delegates to jax.nn.dot_product_attention so the fused backends are used where available.

Source code in xwm/nn/attention.py

def attend(q: Array, k: Array, v: Array, mask: Array | None = None) -> Array:
    """Scaled dot-product attention on ``(heads, N, head_dim)`` inputs.

    ``mask`` is a boolean ``(N_q, N_kv)`` array where ``True`` means *attend*.
    Delegates to :func:`jax.nn.dot_product_attention` so the fused backends are
    used where available.
    """
    q_, k_, v_ = (x.transpose(1, 0, 2)[None] for x in (q, k, v))  # (1, N, H, Dh)
    m = None if mask is None else mask[None, None]  # (1, 1, Nq, Nkv)
    out = jax.nn.dot_product_attention(q_, k_, v_, mask=m)
    return out[0].transpose(1, 0, 2)

block_causal_attention_mask

block_causal_attention_mask(n_blocks: int, block_size: int) -> Array

Causal across blocks, fully connected within a block.

This is the mask an action-conditioned video predictor wants: all tokens of a frame see each other, and frames only see the past.

Source code in xwm/nn/attention.py

def block_causal_attention_mask(n_blocks: int, block_size: int) -> Array:
    """Causal across blocks, fully connected within a block.

    This is the mask an action-conditioned video predictor wants: all tokens of
    a frame see each other, and frames only see the past.
    """
    b = jnp.arange(n_blocks * block_size) // block_size
    return b[:, None] >= b[None, :]

causal_attention_mask

causal_attention_mask(n: int) -> Array

Boolean (n, n) mask allowing each position to see itself and the past.

Source code in xwm/nn/attention.py

def causal_attention_mask(n: int) -> Array:
    """Boolean ``(n, n)`` mask allowing each position to see itself and the past."""
    return jnp.tril(jnp.ones((n, n), dtype=bool))

linear_droppath_schedule

linear_droppath_schedule(depth: int, max_rate: float) -> list[float]

Linearly increasing drop-path rates across depth (the ViT default).

Source code in xwm/nn/drop.py

def linear_droppath_schedule(depth: int, max_rate: float) -> list[float]:
    """Linearly increasing drop-path rates across depth (the ViT default)."""
    if depth == 1:
        return [max_rate]
    return [max_rate * i / (depth - 1) for i in range(depth)]

grid_coords

grid_coords(grid: Grid) -> Array

(prod(grid), len(grid)) integer coordinates, row-major flattened.

Source code in xwm/nn/embed.py

def grid_coords(grid: Grid) -> Array:
    """``(prod(grid), len(grid))`` integer coordinates, row-major flattened."""
    return jnp.asarray(_coords_table(tuple(grid)))

sincos_pos_embed

sincos_pos_embed(grid: Grid, dim: int) -> Array

Fixed sin-cos table for an n-dimensional grid, flattened row-major.

The width is split across axes by axis_dims. For a video grid (T, H, W) this is the factorised 3-D embedding used by V-JEPA.

Returns (prod(grid), dim).

Source code in xwm/nn/embed.py

def sincos_pos_embed(grid: Grid, dim: int) -> Array:
    """Fixed sin-cos table for an n-dimensional grid, flattened row-major.

    The width is split across axes by :func:`axis_dims`. For a video grid
    ``(T, H, W)`` this is the factorised 3-D embedding used by V-JEPA.

    Returns ``(prod(grid), dim)``.
    """
    return jnp.asarray(_sincos_table(tuple(grid), dim))

l2_normalize

l2_normalize(x: Array, axis: int = -1, eps: float = 1e-08) -> Array

Project onto the unit sphere along axis.

Source code in xwm/nn/norm.py

def l2_normalize(x: Array, axis: int = -1, eps: float = 1e-8) -> Array:
    """Project onto the unit sphere along ``axis``."""
    return x / (jnp.linalg.norm(x, axis=axis, keepdims=True) + eps)

unpatchify_2d

unpatchify_2d(tokens: Array, grid: tuple[int, int], patch_size: int, channels: int) -> Array

Fold (N, C * p * p) patch values back into a (C, H, W) image.

Only needed for visualisation and for pixel-space baselines -- predictive world models in xwm never reconstruct pixels during training.

Source code in xwm/nn/patch.py

def unpatchify_2d(tokens: Array, grid: tuple[int, int], patch_size: int, channels: int) -> Array:
    """Fold ``(N, C * p * p)`` patch values back into a ``(C, H, W)`` image.

    Only needed for visualisation and for pixel-space baselines -- predictive
    world models in xwm never reconstruct pixels during training.
    """
    gh, gw = grid
    return rearrange(
        tokens,
        "(gh gw) (c p1 p2) -> c (gh p1) (gw p2)",
        gh=gh,
        gw=gw,
        c=channels,
        p1=patch_size,
        p2=patch_size,
    )

mean_pool

mean_pool(tokens: Array) -> Array

Average over the token axis of an (N, D) sequence.

Source code in xwm/nn/transformer.py

def mean_pool(tokens: Array) -> Array:
    """Average over the token axis of an ``(N, D)`` sequence."""
    return jnp.mean(tokens, axis=0)