RoPE

rope

Rotary Position Embedding (RoPE) primitives.

apply_rotary_emb

apply_rotary_emb(x: array, cos: array, sin: array, *, interleaved: bool = True) -> array

Apply RoPE rotation to x.

Computes x * cos + rotate_half(x) * sin where cos and sin are expanded from the half-dim shape returned by :func:rope_frequencies_1d / :func:rope_frequencies_nd to match the full last-dim of x.

Parameters:

Name	Type	Description	Default
x	array	(..., S, D) tensor to rotate. The trailing two axes must be the sequence and the head dim respectively (broadcasts otherwise).	required
cos	array	(S, D // 2) cosines.	required
sin	array	(S, D // 2) sines.	required
interleaved	bool	See :func:rotate_half. Defaults to True (RoPE-paper convention).	True

Returns:

Type	Description
array	Rotated tensor, same shape as x.

meshgrid_nd

meshgrid_nd(sizes: Sequence[int]) -> list[array]

Build an "ij"-indexed N-D meshgrid of integer positions.

Returns one (S,) flattened position grid per axis, where S = prod(sizes). Use these grids as the position_grids argument to :func:rope_frequencies_nd.

Parameters:

Name	Type	Description	Default
sizes	Sequence[int]	Number of positions along each axis (e.g. [T, H, W] for a video).	required

Returns:

Type	Description
list[array]	List of len(sizes) flattened position arrays, each of
list[array]	shape (prod(sizes),) and dtype float32.

rope_frequencies_1d

rope_frequencies_1d(dim: int, positions: array, theta: float = 10000.0, *, theta_rescale_factor: float = 1.0, interpolation_factor: float = 1.0) -> tuple[array, array]

Compute 1-axis RoPE (cos, sin) for the given positions.

The output has shape (S, dim // 2) and represents the per-pair rotation angle θ_k = pos / theta^(2k/dim) for k ∈ [0, d/2).

Parameters:

Name	Type	Description	Default
dim	int	Head dimension this axis covers (must be even).	required
positions	array	(S,) float array of token positions.	required
theta	float	RoPE base. 10000.0 is the de-facto standard; some recent models use larger bases for longer contexts.	10000.0
theta_rescale_factor	float	NTK-aware rescaling (Reddit user bloc97): theta *= rescale_factor ** (dim / (dim - 2)). Use 1.0 to disable.	1.0
interpolation_factor	float	Position-interpolation factor for context extension (Chen et al., 2023). Positions are scaled by this factor before frequency computation. Use 1.0 to disable.	1.0

Returns:

Type	Description
array	(cos, sin) of shape (S, dim // 2). Pass to
array	func:apply_rotary_emb along with the target tensor.

rope_frequencies_nd

rope_frequencies_nd(dims_per_axis: Sequence[int], position_grids: Sequence[array], theta: float = 10000.0, *, theta_rescale_factor: float | Sequence[float] = 1.0, interpolation_factor: float | Sequence[float] = 1.0) -> tuple[array, array]

Compose multi-axis RoPE by concatenating per-axis frequencies.

For a video transformer with head_dim = sum(dims_per_axis) and each token addressed by a tuple of positions (t, h, w), each axis is RoPE'd independently and the resulting half-angles are concatenated along the last dim.

Parameters:

Name	Type	Description	Default
dims_per_axis	Sequence[int]	Per-axis head-dim allocation, must sum to the attention head dim. All entries must be even.	required
position_grids	Sequence[array]	One (S,) array of positions per axis. S must be the same across axes (one position tuple per token).	required
theta	float	Base frequency, broadcast across axes.	10000.0
theta_rescale_factor	float | Sequence[float]	Either a scalar (broadcast) or one factor per axis. See :func:rope_frequencies_1d.	1.0
interpolation_factor	float | Sequence[float]	Same as above.	1.0

Returns:

Type	Description
tuple[array, array]	(cos, sin) of shape (S, sum(dims_per_axis) // 2).

rotate_half

rotate_half(x: array, *, interleaved: bool = True) -> array

Rotate halves of x for RoPE application.

For interleaved=True (RoPE-paper / Matrix-Game): pairs are adjacent; transforms [x0, x1, x2, x3, ...] into [-x1, x0, -x3, x2, ...].

For interleaved=False (HuggingFace Llama / GPT-NeoX): pairs span halves; transforms [x_left, x_right] into [-x_right, x_left].

Parameters:

Name	Type	Description	Default
x	array	Tensor with even last dim.	required
interleaved	bool	Selects the convention.	True

Returns:

Type	Description
array	Tensor of the same shape with paired elements rotated.

rope

Rotary Position Embeddings (RoPE) — N-axis composable primitives.

RoPE rotates pairs of features in each query/key vector by an angle that depends on the token's position, encoding position into the attention dot product. The math is identical across most models; the practical pitfalls are:

1. Pair layout. Two conventions are in the wild:

2. Interleaved (Llama2 paper, Matrix-Game, RoFormer original): pairs are adjacent — pair k consists of features (x[2k], x[2k+1]).

3. Half-rotated (HuggingFace Transformers default for Llama, GPT-NeoX): pairs span halves — pair k consists of (x[k], x[k + d/2]).

The two compute identical attention scores, but the cos/sin layout and rotate_half form differ. This module supports both via interleaved=True (default, RoPE-paper convention) or interleaved=False.

1. Multi-axis (N-D) composition. Video / image models split the head dim across spatial axes — e.g. CogVideoX uses [t, h, w] with rope_dim_list = [16, 56, 56]. Each axis gets a 1-D RoPE over its own position grid; the resulting per-axis (cos, sin) are concatenated along the feature axis. See :func:rope_frequencies_nd.

2. mx.fast.rope. MLX ships mx.fast.rope for the standard 1-D case but does NOT cover N-axis composition or arbitrary frequency schedules (Megatron non-interleaved, log-spaced SPLIT). This module fills that gap with portable, slow-but-correct primitives.

Model-specific variants (ERNIE-Image Megatron, LTX SPLIT log-spaced) should stay in their ports — arsenal only covers the "standard" case. The shapes returned here mirror the most common HF / Matrix-Game convention.

rope_frequencies_1d

rope_frequencies_1d(dim: int, positions: array, theta: float = 10000.0, *, theta_rescale_factor: float = 1.0, interpolation_factor: float = 1.0) -> tuple[array, array]

Compute 1-axis RoPE (cos, sin) for the given positions.

The output has shape (S, dim // 2) and represents the per-pair rotation angle θ_k = pos / theta^(2k/dim) for k ∈ [0, d/2).

Parameters:

Name	Type	Description	Default
dim	int	Head dimension this axis covers (must be even).	required
positions	array	(S,) float array of token positions.	required
theta	float	RoPE base. 10000.0 is the de-facto standard; some recent models use larger bases for longer contexts.	10000.0
theta_rescale_factor	float	NTK-aware rescaling (Reddit user bloc97): theta *= rescale_factor ** (dim / (dim - 2)). Use 1.0 to disable.	1.0
interpolation_factor	float	Position-interpolation factor for context extension (Chen et al., 2023). Positions are scaled by this factor before frequency computation. Use 1.0 to disable.	1.0

Returns:

Type	Description
array	(cos, sin) of shape (S, dim // 2). Pass to
array	func:apply_rotary_emb along with the target tensor.

rope_frequencies_nd

rope_frequencies_nd(dims_per_axis: Sequence[int], position_grids: Sequence[array], theta: float = 10000.0, *, theta_rescale_factor: float | Sequence[float] = 1.0, interpolation_factor: float | Sequence[float] = 1.0) -> tuple[array, array]

Compose multi-axis RoPE by concatenating per-axis frequencies.

For a video transformer with head_dim = sum(dims_per_axis) and each token addressed by a tuple of positions (t, h, w), each axis is RoPE'd independently and the resulting half-angles are concatenated along the last dim.

Parameters:

Name	Type	Description	Default
dims_per_axis	Sequence[int]	Per-axis head-dim allocation, must sum to the attention head dim. All entries must be even.	required
position_grids	Sequence[array]	One (S,) array of positions per axis. S must be the same across axes (one position tuple per token).	required
theta	float	Base frequency, broadcast across axes.	10000.0
theta_rescale_factor	float | Sequence[float]	Either a scalar (broadcast) or one factor per axis. See :func:rope_frequencies_1d.	1.0
interpolation_factor	float | Sequence[float]	Same as above.	1.0

Returns:

Type	Description
tuple[array, array]	(cos, sin) of shape (S, sum(dims_per_axis) // 2).

rotate_half

rotate_half(x: array, *, interleaved: bool = True) -> array

Rotate halves of x for RoPE application.

For interleaved=True (RoPE-paper / Matrix-Game): pairs are adjacent; transforms [x0, x1, x2, x3, ...] into [-x1, x0, -x3, x2, ...].

For interleaved=False (HuggingFace Llama / GPT-NeoX): pairs span halves; transforms [x_left, x_right] into [-x_right, x_left].

Parameters:

Name	Type	Description	Default
x	array	Tensor with even last dim.	required
interleaved	bool	Selects the convention.	True

Returns:

Type	Description
array	Tensor of the same shape with paired elements rotated.

apply_rotary_emb

apply_rotary_emb(x: array, cos: array, sin: array, *, interleaved: bool = True) -> array

Apply RoPE rotation to x.

Computes x * cos + rotate_half(x) * sin where cos and sin are expanded from the half-dim shape returned by :func:rope_frequencies_1d / :func:rope_frequencies_nd to match the full last-dim of x.

Parameters:

Name	Type	Description	Default
x	array	(..., S, D) tensor to rotate. The trailing two axes must be the sequence and the head dim respectively (broadcasts otherwise).	required
cos	array	(S, D // 2) cosines.	required
sin	array	(S, D // 2) sines.	required
interleaved	bool	See :func:rotate_half. Defaults to True (RoPE-paper convention).	True

Returns:

Type	Description
array	Rotated tensor, same shape as x.

meshgrid_nd

meshgrid_nd(sizes: Sequence[int]) -> list[array]

Build an "ij"-indexed N-D meshgrid of integer positions.

Returns one (S,) flattened position grid per axis, where S = prod(sizes). Use these grids as the position_grids argument to :func:rope_frequencies_nd.

Parameters:

Name	Type	Description	Default
sizes	Sequence[int]	Number of positions along each axis (e.g. [T, H, W] for a video).	required

Returns:

Type	Description
list[array]	List of len(sizes) flattened position arrays, each of
list[array]	shape (prod(sizes),) and dtype float32.