From d12124ce5bfc3165448514284a69c1030dab936d Mon Sep 17 00:00:00 2001 From: Carlos Trujillo Date: Fri, 7 Aug 2026 23:27:21 +0300 Subject: [PATCH 1/2] Add RMSNorm and rotary position embeddings Two additive decoder primitives, from the in-scope set agreed on #53. Both are `UnaryLayerOp`s with complete `__props__`, so they serialize as leaves for free, and neither touches the `__props__` of an already-serialized op -- no `GRAPH_FORMAT_VERSION` bump. RMSNorm (`RMSNorm` / `RMSNormLayer`, beside `LayerNorm`): scale-only affine with no shift and no mean subtraction, matching `torch.nn.RMSNorm`, `flax.linen/nnx.RMSNorm` and `tinygrad.nn.RMSNorm`. Default epsilon is 1e-6 (flax/tinygrad), deliberately not `LayerNorm`'s 1e-5 (torch). RoPE (`rotary_embedding` / `RotaryEmbedding` / `RotaryEmbeddingLayer` in a new `layers/positional.py`): positions are an explicit input rather than an implied `0..seq-1`, so one graph serves a full sequence and a single decode step. Both pairing conventions are props, because they are not interchangeable and the choice is fixed by the weights: `"half"` pairs `i` with `i + d/2` (HuggingFace `rotate_half`, GPT-NeoX, flax) and `"adjacent"` pairs `2i` with `2i + 1` (RoFormer, GPT-J, torchtune). Linear position interpolation and static NTK-aware scaling are also props. Composes with the existing `AttentionLayer` unchanged -- `_sdpa_graph` is already position-agnostic. `_constant_like` moves from `activations.py` to `base.py` as `constant_like`: both the activations and the new norm need dtype-pinned scalar constants, and `base.py` is already the shared root neither can cycle through. Without it an epsilon of 1e-6 widens a float32 input to float64, since pytensor's autocaster types a bare Python float by value. Correctness is pinned against an independent scalar reference in `tests/test_positional.py` that imports neither pytensor nor pytensor_ml and applies each 2x2 rotation literally: flipping the pairing on the reference side fails the test. Also covered: one-step-at-a-time decoding equals the whole-sequence result, scores depend only on relative position, rotation is orthogonal, position 0 is the identity, and linear scaling by 4 at position 8 equals plain RoPE at position 2. Both ops lower on JAX through pytensor's `OpFromGraph` fallback with no `jax_funcify` of their own, asserted in `tests/dispatch/jax/test_positional.py`. The interleaved pairing therefore reassembles its output with strided `set_subtensor` writes rather than `split_dims`/`join_dims`, which have no JAX or MLX conversion. --- pytensor_ml/activations.py | 32 +-- pytensor_ml/base.py | 21 +- pytensor_ml/layers/__init__.py | 10 + pytensor_ml/layers/norm.py | 122 +++++++++- pytensor_ml/layers/positional.py | 335 ++++++++++++++++++++++++++ tests/dispatch/jax/test_positional.py | 45 ++++ tests/test_layers.py | 79 +++++- tests/test_positional.py | 329 +++++++++++++++++++++++++ tests/test_serialize.py | 26 ++ 9 files changed, 970 insertions(+), 29 deletions(-) create mode 100644 pytensor_ml/layers/positional.py create mode 100644 tests/dispatch/jax/test_positional.py create mode 100644 tests/test_positional.py diff --git a/pytensor_ml/activations.py b/pytensor_ml/activations.py index f1cf253..186ad52 100644 --- a/pytensor_ml/activations.py +++ b/pytensor_ml/activations.py @@ -1,23 +1,7 @@ import numpy as np import pytensor.tensor as pt -from pytensor import config - -from pytensor_ml.base import Layer - - -def _constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable: - """ - Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype. - - PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand - depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x`` - promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence. - """ - dtype = np.dtype(x.dtype) - # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision. - dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX) - return pt.constant(np.asarray(value, dtype=dtype)) +from pytensor_ml.base import Layer, constant_like class Activation(Layer): ... @@ -62,7 +46,7 @@ def __init__(self, negative_slope: float = 0.01): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - out = pt.switch(x > 0, x, _constant_like(self.negative_slope, x) * x) + out = pt.switch(x > 0, x, constant_like(self.negative_slope, x) * x) out.name = "LeakyReLU" return out @@ -149,14 +133,14 @@ def __init__(self, approximate: bool = True): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - half = _constant_like(0.5, x) - one = _constant_like(1.0, x) + half = constant_like(0.5, x) + one = constant_like(1.0, x) if self.approximate: - sqrt_2_over_pi = _constant_like(np.sqrt(2.0 / np.pi), x) - cubic_coef = _constant_like(0.044715, x) + sqrt_2_over_pi = constant_like(np.sqrt(2.0 / np.pi), x) + cubic_coef = constant_like(0.044715, x) out = half * x * (one + pt.tanh(sqrt_2_over_pi * (x + cubic_coef * x**3))) else: - sqrt2 = _constant_like(np.sqrt(2.0), x) + sqrt2 = constant_like(np.sqrt(2.0), x) out = half * x * (one + pt.erf(x / sqrt2)) out.name = "GELU" return out @@ -182,7 +166,7 @@ def __init__(self, beta: float = 1.0): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - out = x * pt.sigmoid(_constant_like(self.beta, x) * x) + out = x * pt.sigmoid(constant_like(self.beta, x) * x) out.name = "Swish" return out diff --git a/pytensor_ml/base.py b/pytensor_ml/base.py index c8baf40..42a6864 100644 --- a/pytensor_ml/base.py +++ b/pytensor_ml/base.py @@ -1,12 +1,31 @@ from abc import ABC, abstractmethod from typing import Protocol, runtime_checkable +import numpy as np import pytensor.tensor as pt +from pytensor import config from pytensor.compile.builders import SymbolicOp from pytensor.tensor.variable import TensorVariable +def constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable: + """ + Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype. + + PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand + depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x`` + promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence. + + Lives here, beside the layer base classes, because both ``pytensor_ml.activations`` and the layer + modules need it and this module is already the shared root that neither can cycle through. + """ + dtype = np.dtype(x.dtype) + # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision. + dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX) + return pt.constant(np.asarray(value, dtype=dtype)) + + class Layer(ABC): """Base class for the objects that build layer graphs. Defined here, not in ``pytensor_ml.layers``, so that ``pytensor_ml.activations`` can subclass it without a circular import.""" @@ -53,4 +72,4 @@ def update_map(self) -> dict[int, int]: """Map each output index to the index of the input that output updates.""" -__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp"] +__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp", "constant_like"] diff --git a/pytensor_ml/layers/__init__.py b/pytensor_ml/layers/__init__.py index ba5b128..3451217 100644 --- a/pytensor_ml/layers/__init__.py +++ b/pytensor_ml/layers/__init__.py @@ -18,6 +18,13 @@ LayerNormLayer, NoRunningStatsBatchNormLayer, PredictionBatchNormLayer, + RMSNorm, + RMSNormLayer, +) +from pytensor_ml.layers.positional import ( + RotaryEmbedding, + RotaryEmbeddingLayer, + rotary_embedding, ) from pytensor_ml.layers.transformer import FeedForward, TransformerBlock @@ -33,8 +40,11 @@ "LayerNorm", "Linear", "MultiheadAttention", + "RMSNorm", + "RotaryEmbedding", "Sequential", "Squeeze", "TransformerBlock", + "rotary_embedding", "scaled_dot_product_attention", ] diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py index ed264ca..8c688ce 100644 --- a/pytensor_ml/layers/norm.py +++ b/pytensor_ml/layers/norm.py @@ -3,7 +3,7 @@ from pytensor import config -from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp +from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp, constant_like from pytensor_ml.params import NonTrainableParameter, TrainableParameter, non_trainable, trainable @@ -31,6 +31,21 @@ def _standardize(X, epsilon, axis, keepdims=False): return (X - mu) / pt.sqrt(sigma_sq + epsilon), mu, sigma_sq +def _rms_normalize(X, epsilon): + """ + Scale ``X`` to unit root mean square over its last axis, without centering it. + + Notes + ----- + RMSNorm divides by the root mean square rather than the standard deviation, so unlike + :func:`_standardize` it leaves the mean of ``X`` intact. That omission is the entire difference + between the two, and pretrained weights depend on it -- do not "simplify" this into a + :func:`_standardize` call. + """ + mean_square = pt.mean(pt.square(X), axis=-1, keepdims=True) + return X / pt.sqrt(mean_square + constant_like(epsilon, X)) + + def _affine_input_count(affine: bool) -> int: """Number of inputs the learned affine transform contributes, which every norm op places directly after ``X``. Both the graph builders and :meth:`BatchNormLayer.update_map` index around it.""" @@ -75,12 +90,17 @@ def _resolve_n_in(name: str, n_in: int | None, X: pt.TensorVariable | None) -> i return inferred +def _scale_parameter(name: str, n_in: int) -> TrainableParameter: + """Build the learned scale. RMSNorm has only this one; the shift-and-scale norms pair it with a + ``loc``, so the ``_scale`` suffix is fixed here and nowhere else.""" + return trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale") + + def _affine_parameters(name: str, n_in: int) -> tuple[TrainableParameter, TrainableParameter]: """Build the learned shift and scale. Returns them in the ``(loc, scale)`` order that every norm op unpacks its inputs in, so the two cannot drift apart.""" loc = trainable(np.zeros(n_in, dtype=config.floatX), f"{name}_loc") - scale = trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale") - return loc, scale + return loc, _scale_parameter(name, n_in) class BatchNormLayer(LayerOp): @@ -335,3 +355,99 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable: X_transformed.name = f"{self.name}_output" return X_transformed + + +class RMSNormLayer(UnaryLayerOp): + __props__ = ("n_in", "epsilon", "affine") + + def build_inner_graph(self, X, *rest): + X_normalized = _rms_normalize(X, self.epsilon) + if not self.affine: + return [X_normalized] + + # Scale-only, so the affine transform contributes one input rather than the ``(loc, scale)`` + # pair the other norm ops unpack; the single-element unpack asserts that arity. + (scale,) = rest + return [X_normalized * scale] + + +class RMSNorm(Layer): + r""" + Root-mean-square layer normalization over the last (feature) axis. + + Divide each sample by the root mean square of its own features, then optionally apply a learned + scale: + + .. math:: + + y = \frac{x}{\sqrt{\frac{1}{n} \sum_i x_i^2 + \epsilon}} \cdot \gamma. + + Unlike :class:`LayerNorm` there is no mean subtraction and no learned shift. That is not a + simplification of this implementation but the definition: :class:`torch.nn.RMSNorm`, + ``flax.linen.RMSNorm``/``flax.nnx.RMSNorm`` and ``tinygrad.nn.RMSNorm`` all expose a weight and + no bias. It is the normalization used by the Llama, Gemma and Qwen decoder families. + + Parameters + ---------- + name : str, optional + Name used as a prefix for the layer's parameters. Default is "RMSNorm". + n_in : int, optional + Size of the normalized feature axis. Inferred from the input's last dimension on the first + call when omitted. + epsilon : float, optional + Constant :math:`\epsilon` added to the mean square for numerical stability. Default is 1e-6, + matching ``flax``'s and ``tinygrad``'s RMSNorm. This deliberately differs from + :class:`LayerNorm`'s 1e-5, which follows :class:`torch.nn.LayerNorm`; pretrained weights are + published against their own framework's value. + affine : bool, optional + Apply the learned scale :math:`\gamma`. There is no shift to disable. Default is True. + """ + + def __init__( + self, + name: str | None = None, + n_in: int | None = None, + epsilon: float = 1e-6, + affine: bool = True, + ): + self.name = name if name else "RMSNorm" + self.n_in = n_in + self.epsilon = epsilon + self.affine = affine + + self.scale: TrainableParameter | None = None + + self.initialized = False + self._initialize_params(None) + + def _initialize_params(self, X: pt.TensorVariable | None): + if self.initialized: + return + + n_in = _resolve_n_in(self.name, self.n_in, X) + if n_in is None: + return + + if self.affine: + self.scale = _scale_parameter(self.name, n_in) + + self.initialized = True + + def __call__(self, X: pt.TensorLike) -> pt.TensorVariable: + X = pt.as_tensor(X) + self._initialize_params(X) + + inputs = [X] + if self.affine: + assert self.scale is not None + inputs.append(self.scale) + + X_transformed = RMSNormLayer( + name=self.name, + n_in=self.n_in, + epsilon=self.epsilon, + affine=self.affine, + )(*inputs) + X_transformed.name = f"{self.name}_output" + + return X_transformed diff --git a/pytensor_ml/layers/positional.py b/pytensor_ml/layers/positional.py new file mode 100644 index 0000000..5110075 --- /dev/null +++ b/pytensor_ml/layers/positional.py @@ -0,0 +1,335 @@ +from typing import Literal, get_args + +import numpy as np +import pytensor.tensor as pt + +from pytensor.tensor.variable import TensorVariable + +from pytensor_ml.base import Layer, UnaryLayerOp + +Pairing = Literal["half", "adjacent"] +Scaling = Literal["none", "linear", "ntk"] + +_PAIRINGS = get_args(Pairing) +_SCALINGS = get_args(Scaling) + + +def _validate_options(pairing: str, scaling: str, scaling_factor: float) -> None: + """Reject unknown options where they are given, rather than letting a typo silently select the + default behaviour inside the inner graph.""" + if pairing not in _PAIRINGS: + raise ValueError(f"pairing must be one of {_PAIRINGS}, got {pairing!r}") + if scaling not in _SCALINGS: + raise ValueError(f"scaling must be one of {_SCALINGS}, got {scaling!r}") + if scaling != "none" and scaling_factor <= 0: + raise ValueError(f"scaling_factor must be positive, got {scaling_factor}") + + +def _resolve_head_dim(x: TensorVariable) -> int: + """The rotated feature size, which must be static and even. + + Static because the frequencies are materialized as a constant at graph-build time, and even + because the rotation acts on two-dimensional subspaces. + """ + if not np.issubdtype(np.dtype(x.dtype), np.floating): + raise ValueError( + f"RotaryEmbedding rotates a floating-point tensor, but got dtype {x.dtype}. Cast the " + f"input first." + ) + + head_dim = x.type.shape[-1] + if head_dim is None: + raise ValueError( + "RotaryEmbedding needs a statically known last dimension to build its frequencies. " + "Give the input a static feature dimension, e.g. pt.tensor(shape=(None, n_head, None, " + "head_dim))." + ) + if head_dim % 2: + raise ValueError(f"RotaryEmbedding needs an even head dimension, got {head_dim}.") + + return head_dim + + +def _align_to_head_axes( + angles: TensorVariable, x: TensorVariable, position_ids: TensorVariable +) -> TensorVariable: + """Insert ``x``'s head axes into the angle tensor so the sequence axes line up. + + ``position_ids`` describes tokens: its last axis is the sequence and any leading axes are batch + axes. ``x`` carries its head axes *between* those two, so the angles need matching placeholder + axes. How many follows from the ranks, which is why the caller does not pass an unsqueeze position: + ``(seq,)`` broadcasts over everything, and ``(batch, seq)`` lines up against + ``(batch, n_head, seq, head_dim)`` without the caller reshaping. Deriving it also keeps the + inserted axes statically 1, which is what pytensor requires before it will broadcast them -- a + hand-built ``(batch, 1, seq)`` input whose middle axis is merely unknown is rejected at runtime. + """ + n_head_axes = (x.type.ndim - 1) - position_ids.type.ndim + if n_head_axes < 0: + raise ValueError( + f"position_ids has {position_ids.type.ndim} dimensions, more than the " + f"{x.type.ndim - 1} non-feature dimensions of x. Its last axis is the sequence and any " + f"leading axes are batch axes, so it can never be wider than x minus the feature axis." + ) + if n_head_axes == 0: + return angles + + return angles[(Ellipsis, *(None,) * n_head_axes, slice(None), slice(None))] + + +def _inverse_frequencies( + head_dim: int, base: float, scaling: str, scaling_factor: float +) -> np.ndarray: + r""" + Angular frequencies :math:`\theta_i = \mathrm{base}^{-2i/d}` for each rotated pair. + + Computed in numpy at graph-build time: ``base``, ``head_dim`` and the scaling are all known then, + so the frequencies enter the graph as one constant instead of a subgraph of ``arange`` and + ``power``. + + Parameters + ---------- + head_dim : int + Size of the rotated feature axis. There are ``head_dim // 2`` frequencies. + base : float + Geometric base of the frequency ladder, :math:`\theta` in the RoFormer paper. + scaling : str + ``"none"``, ``"linear"`` for position interpolation, or ``"ntk"`` for NTK-aware scaling. + scaling_factor : float + Context-extension factor for the scaled variants; ignored when ``scaling="none"``. + + Returns + ------- + ndarray + Frequencies of shape ``(head_dim // 2,)``, in float64. + """ + if scaling == "ntk": + if head_dim <= 2: + raise ValueError( + f"NTK-aware scaling rescales the base by scaling_factor ** (d / (d - 2)), which is " + f"undefined for head_dim <= 2; got {head_dim}." + ) + # Static NTK-aware scaling: stretch the frequency ladder itself rather than the positions, so + # the highest frequencies are left almost untouched. This is the exponent HuggingFace's + # `_compute_dynamic_ntk_parameters` uses; the dynamic form, which recomputes the base from the + # running sequence length, needs a runtime length and is deliberately not offered here. + base = base * scaling_factor ** (head_dim / (head_dim - 2)) + + inverse_frequencies = 1.0 / base ** (np.arange(0, head_dim, 2, dtype="float64") / head_dim) + + if scaling == "linear": + # Position interpolation: dividing the frequencies is algebraically the same as dividing the + # positions, because the angle is bilinear in the two. + inverse_frequencies = inverse_frequencies / scaling_factor + + return inverse_frequencies + + +def _split_pairs( + x: TensorVariable, pairing: str, half: int +) -> tuple[TensorVariable, TensorVariable]: + """Split the feature axis into the two halves of every rotated pair. + + The two conventions differ only in which channels are paired, but they are not interchangeable: + weights trained under one produce nonsense under the other. + + ``"half"`` pairs channel ``i`` with ``i + d/2``, matching HuggingFace's ``rotate_half`` (Llama, + Qwen, Mistral, Gemma), GPT-NeoX and flax's ``RoPE``. ``"adjacent"`` pairs ``2i`` with ``2i + 1``, + matching the original RoFormer formulation, GPT-J and ``torchtune``. + """ + if pairing == "half": + return x[..., :half], x[..., half:] + + return x[..., 0::2], x[..., 1::2] + + +def _join_pairs( + x: TensorVariable, first: TensorVariable, second: TensorVariable, pairing: str +) -> TensorVariable: + """Reassemble the feature axis, inverting :func:`_split_pairs` for the same ``pairing``. + + Both branches stay on ops every backend lowers, and both preserve the static feature size that + lets the backends vectorize the rotation. Notably ``split_dims``/``join_dims`` would read more + naturally for the interleaved case but have no JAX or MLX conversion, so a graph using them + compiles only in the default backend. + """ + if pairing == "half": + return pt.concatenate([first, second], axis=-1) + + # Strided writes into a copy of x rather than an interleaving reshape. Every channel is + # overwritten -- the even ones by `first`, the odd ones by `second` -- so no value of x survives, + # and both were read before either write. + rotated = pt.set_subtensor(x[..., 0::2], first) + return pt.set_subtensor(rotated[..., 1::2], second) + + +class RotaryEmbeddingLayer(UnaryLayerOp): + __props__ = ("base", "pairing", "scaling", "scaling_factor") + + def build_inner_graph(self, x, position_ids): + _validate_options(self.pairing, self.scaling, self.scaling_factor) + head_dim = _resolve_head_dim(x) + + inverse_frequencies = _inverse_frequencies( + head_dim, self.base, self.scaling, self.scaling_factor + ) + # Angles, and therefore cos/sin, are computed in x's own dtype: floatX is float32 or float64 + # here, so this already is the "compute in fp32" path that the fp16/bf16 frameworks upcast to. + angles = position_ids[..., None].astype(x.dtype) * pt.constant( + inverse_frequencies.astype(x.dtype), name="inverse_frequencies" + ) + angles = _align_to_head_axes(angles, x, position_ids) + cos, sin = pt.cos(angles), pt.sin(angles) + + first, second = _split_pairs(x, self.pairing, head_dim // 2) + rotated = _join_pairs( + x, first * cos - second * sin, second * cos + first * sin, self.pairing + ) + + return [rotated] + + +def rotary_embedding( + x: pt.TensorLike, + position_ids: pt.TensorLike, + *, + base: float = 10_000.0, + pairing: Pairing = "half", + scaling: Scaling = "none", + scaling_factor: float = 1.0, +) -> TensorVariable: + r""" + Rotary position embedding (RoPE) applied to the trailing feature axis. + + Rotate each two-dimensional subspace of the feature axis by an angle proportional to the token's + position: + + .. math:: + + \begin{pmatrix} x'_a \\ x'_b \end{pmatrix} = + \begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ + \sin m\theta_i & \cos m\theta_i \end{pmatrix} + \begin{pmatrix} x_a \\ x_b \end{pmatrix}, + \qquad \theta_i = \mathrm{base}^{-2i/d}, + + where :math:`m` is the position and :math:`(a, b)` is the ``i``-th channel pair. Because a + rotation is orthogonal, the dot product of a rotated query and a rotated key depends only on the + *difference* of their positions -- that relative-position property is the point of the scheme, and + it is what makes RoPE usable for incremental decoding: a token rotated by its absolute position + keeps the same relationship to every earlier token no matter when it was computed. + + Apply this to queries and keys before :func:`~pytensor_ml.layers.attention.scaled_dot_product_attention`, + never to values. + + Parameters + ---------- + x : TensorLike + Tensor whose last axis is rotated, typically queries or keys of shape + ``(..., n_head, seq, head_dim)``. ``head_dim`` must be static and even. + position_ids : TensorLike + Position of each token, shape ``(..., seq)``: the last axis is the sequence and any leading + axes are batch axes. ``x``'s head axes are inserted for you, so both ``(seq,)`` -- shared by + the whole batch -- and ``(batch, seq)`` work against ``(batch, n_head, seq, head_dim)``. + Positions are explicit rather than an implied ``0..seq-1`` so that the same graph serves a full + sequence and a single decode step appended to a cached prefix. + base : float, optional + Geometric base of the frequency ladder (:math:`\theta` in some formulations). Default is + 10000.0, the value used by RoFormer, Llama, GPT-NeoX and torchtune. + pairing : str, optional + Which channels form each rotated pair. ``"half"`` (default) pairs ``i`` with ``i + d/2``, as + HuggingFace's ``rotate_half``, GPT-NeoX and flax's ``RoPE`` do; ``"adjacent"`` pairs ``2i`` + with ``2i + 1``, as the original RoFormer, GPT-J and torchtune do. The two are a permutation of + the feature axis apart, so this must match whatever the weights were trained with. + scaling : str, optional + Context-extension scheme. ``"none"`` (default) for plain RoPE, ``"linear"`` for position + interpolation (Chen et al. 2023), or ``"ntk"`` for static NTK-aware scaling, which rescales the + base by ``scaling_factor ** (head_dim / (head_dim - 2))``. + scaling_factor : float, optional + Extension factor for the scaled variants, ignored when ``scaling="none"``. Default is 1.0, + which is the identity for both variants. + + Returns + ------- + TensorVariable + ``x`` with its last axis rotated, same shape and dtype. + """ + _validate_options(pairing, scaling, scaling_factor) + + x = pt.as_tensor(x) + position_ids = pt.as_tensor(position_ids) + + result = RotaryEmbeddingLayer( + name="RotaryEmbedding", + base=base, + pairing=pairing, + scaling=scaling, + scaling_factor=scaling_factor, + )(x, position_ids) + result.name = "rotary_embedding_output" + + return result + + +class RotaryEmbedding(Layer): + r""" + Rotary position embeddings as a reusable, configured layer. + + Holds the frequency configuration so that queries and keys are rotated identically -- they must be, + since attention compares them -- and calls :func:`rotary_embedding` to build the graph. The layer + has no parameters of its own. + + Parameters + ---------- + name : str or None + Name prefix for the layer's output. Defaults to "RotaryEmbedding" when None. + base : float, optional + Geometric base of the frequency ladder. Default is 10000.0. + pairing : str, optional + ``"half"`` or ``"adjacent"``. Default is ``"half"``. + scaling : str, optional + ``"none"``, ``"linear"``, or ``"ntk"``. Default is ``"none"``. + scaling_factor : float, optional + Extension factor for the scaled variants. Default is 1.0. + + See Also + -------- + rotary_embedding : The functional form, documenting the conventions in full. + """ + + def __init__( + self, + name: str | None = None, + base: float = 10_000.0, + pairing: Pairing = "half", + scaling: Scaling = "none", + scaling_factor: float = 1.0, + ): + _validate_options(pairing, scaling, scaling_factor) + + self.name = name if name else "RotaryEmbedding" + self.base = base + self.pairing = pairing + self.scaling = scaling + self.scaling_factor = scaling_factor + + # Positions are a required second input, which the one-tensor Layer.__call__ signature does not + # describe. Widening the base class would loosen it for every layer that really is unary. + def __call__( # type: ignore[override] + self, x: pt.TensorLike, position_ids: pt.TensorLike + ) -> TensorVariable: + out = rotary_embedding( + x, + position_ids, + base=self.base, + pairing=self.pairing, + scaling=self.scaling, + scaling_factor=self.scaling_factor, + ) + out.name = f"{self.name}_output" + + return out + + +__all__ = [ + "RotaryEmbedding", + "rotary_embedding", +] diff --git a/tests/dispatch/jax/test_positional.py b/tests/dispatch/jax/test_positional.py new file mode 100644 index 0000000..df2e453 --- /dev/null +++ b/tests/dispatch/jax/test_positional.py @@ -0,0 +1,45 @@ +from functools import partial + +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +pytest.importorskip("jax") + +from pytensor_ml.layers.norm import RMSNorm +from pytensor_ml.layers.positional import rotary_embedding +from tests.dispatch.jax.test_basic import compare_jax_and_py + +floatX = pytensor.config.floatX +assert_close = partial(np.testing.assert_allclose, atol=1e-5, rtol=1e-4) + + +@pytest.fixture(scope="module") +def rng(): + return np.random.default_rng(sum(map(ord, "JAX Positional"))) + + +# Neither op registers a jax_funcify of its own: both are SymbolicOps whose inner graphs are built +# from ops the backend already knows, so pytensor's OpFromGraph fallback inlines them. These tests +# hold that claim to account -- a primitive the backend cannot lower is not a usable primitive. +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_rotary_embedding_matches_py(pairing, rng): + x_np = rng.normal(size=(2, 3, 5, 8)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + + out = rotary_embedding(x, positions, pairing=pairing) + compare_jax_and_py([x, positions], out, [x_np, np.arange(5)], assert_fn=assert_close) + + +@pytest.mark.parametrize("affine", [True, False], ids=["affine", "no_affine"]) +def test_rms_norm_matches_py(affine, rng): + X_np = rng.normal(size=(8, 6)).astype(floatX) + X = pt.tensor("X", shape=(None, 6)) + + rms_norm = RMSNorm("rms", n_in=6, affine=affine) + if affine: + rms_norm.scale.set_value(rng.normal(size=6).astype(floatX)) + + compare_jax_and_py([X], rms_norm(X), [X_np], assert_fn=assert_close) diff --git a/tests/test_layers.py b/tests/test_layers.py index f58e50a..90e6907 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -10,7 +10,16 @@ import pytensor_ml.layers from pytensor_ml.activations import ReLU -from pytensor_ml.layers import BatchNorm2D, Dropout, Embedding, Input, LayerNorm, Linear, Sequential +from pytensor_ml.layers import ( + BatchNorm2D, + Dropout, + Embedding, + Input, + LayerNorm, + Linear, + RMSNorm, + Sequential, +) from pytensor_ml.pytensorf import ( collect_non_trainable_updates, collect_trainable_params, @@ -220,6 +229,72 @@ def test_layer_norm_no_affine_standardizes_each_row(rng): np.testing.assert_allclose(res.var(axis=-1), 1.0, rtol=1e-3) +@pytest.mark.parametrize("batch_shape", [(10,), (2, 4)], ids=["2d", "3d"]) +@pytest.mark.parametrize("n_in", [6, None], ids=["specified", "lazy"]) +def test_rms_norm_forward(n_in, batch_shape, rng): + X = pt.tensor("X", shape=(*(None,) * len(batch_shape), 6)) + rms_norm = RMSNorm(name="RMSNorm_1", n_in=n_in) + out = rms_norm(X) + assert out.name == "RMSNorm_1_output" + + X_np = rng.normal(size=(*batch_shape, 6)).astype(floatX) + scale_np = rng.normal(size=(6,)).astype(floatX) + rms_norm.scale.set_value(scale_np) + + res = out.eval({X: X_np}) + mean_square_np = np.square(X_np).mean(axis=-1, keepdims=True) + expected = X_np / np.sqrt(mean_square_np + rms_norm.epsilon) * scale_np + np.testing.assert_allclose(res, expected, rtol=1e-5) + + +def test_rms_norm_has_no_shift_parameter(): + # Scale-only by definition, matching torch.nn.RMSNorm and flax/tinygrad's RMSNorm. A `loc` here + # would mean the layer had quietly become a LayerNorm. + rms_norm = RMSNorm(name="RMSNorm_1", n_in=6) + assert not hasattr(rms_norm, "loc") + assert collect_trainable_params(rms_norm(pt.tensor("X", shape=(None, 6)))) == [rms_norm.scale] + + +def test_rms_norm_no_affine_gives_unit_root_mean_square(rng): + X = pt.tensor("X", shape=(None, 8)) + out = RMSNorm(name="RMSNorm_1", n_in=8, affine=False)(X) + + X_np = rng.normal(loc=3.0, scale=2.0, size=(10, 8)).astype(floatX) + res = out.eval({X: X_np}) + + np.testing.assert_allclose(np.sqrt(np.square(res).mean(axis=-1)), 1.0, rtol=1e-3) + + +def test_rms_norm_does_not_centre_its_input(rng): + # The one thing that separates RMSNorm from LayerNorm: the mean survives. On input with a large + # offset, LayerNorm removes it and RMSNorm does not, so the two must disagree. + X = pt.tensor("X", shape=(None, 8)) + X_np = rng.normal(loc=5.0, scale=1.0, size=(10, 8)).astype(floatX) + + rms = RMSNorm(name="rms", n_in=8, affine=False)(X).eval({X: X_np}) + layer = LayerNorm(name="ln", n_in=8, affine=False)(X).eval({X: X_np}) + + # LayerNorm removes the offset outright; RMSNorm only rescales it, so a clearly non-zero mean + # survives. A stray mean subtraction in RMSNorm would drive this to zero. + np.testing.assert_allclose(layer.mean(axis=-1), 0.0, atol=1e-5) + assert np.all(rms.mean(axis=-1) > 0.5) + assert not np.allclose(rms, layer) + + +def test_rms_norm_prediction_matches_training(rng): + # Per-sample statistics, identical in train and eval, so like LayerNorm it needs no prediction + # rewrite: rewrite_for_prediction must leave its output unchanged. + X = pt.tensor("X", shape=(None, 6)) + rms_norm = RMSNorm("rms", n_in=6) + out = rms_norm(X) + rms_norm.scale.set_value(rng.normal(size=6).astype(floatX)) + + X_np = rng.normal(size=(10, 6)).astype(floatX) + np.testing.assert_allclose( + rewrite_for_prediction(out).eval({X: X_np}), out.eval({X: X_np}), rtol=1e-6 + ) + + def test_batch_norm_2d_learns_population_stats(rng): population_mean, population_std = 3.2, 6.2 X = pt.tensor("X", shape=(None, 32)) @@ -285,6 +360,8 @@ def sample_batch(): ("NoRunningStatsBatchNormLayer", "norm"), ("PredictionBatchNormLayer", "norm"), ("LayerNormLayer", "norm"), + ("RMSNormLayer", "norm"), + ("RotaryEmbeddingLayer", "positional"), ], ) def test_marker_ops_stay_reachable_from_the_package(op_name, submodule): diff --git a/tests/test_positional.py b/tests/test_positional.py new file mode 100644 index 0000000..2876043 --- /dev/null +++ b/tests/test_positional.py @@ -0,0 +1,329 @@ +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +from pytensor_ml.layers.attention import scaled_dot_product_attention +from pytensor_ml.layers.positional import RotaryEmbedding, rotary_embedding +from pytensor_ml.pytensorf import rewrite_for_prediction + +floatX = pytensor.config.floatX + + +@pytest.fixture(scope="module") +def rng(): + return np.random.default_rng(sum(map(ord, "Rotary Test"))) + + +def rope_np(x, positions, *, base=10_000.0, pairing="half", scaling="none", scaling_factor=1.0): + """ + Independent RoPE reference, built pair by pair from explicit 2x2 rotations. + + Deliberately written from the RoFormer definition rather than from the implementation under test: + it imports no pytensor and no pytensor_ml, indexes scalars instead of slicing halves, and applies + the rotation matrix literally. A shared mistake in the pairing, a sign, or the frequency ladder + therefore cannot cancel out between the two. + + Parameters + ---------- + x : ndarray + Shape ``(..., seq, head_dim)``. + positions : ndarray + Shape ``(seq,)``, shared by every leading axis of ``x``. + """ + x = np.asarray(x, dtype="float64") + positions = np.asarray(positions) + head_dim = x.shape[-1] + half = head_dim // 2 + + if scaling == "ntk": + base = base * scaling_factor ** (head_dim / (head_dim - 2)) + + theta = np.array([base ** (-2.0 * i / head_dim) for i in range(half)]) + if scaling == "linear": + theta = theta / scaling_factor + + out = np.empty_like(x) + for lead in np.ndindex(*x.shape[:-2]): + for step, position in enumerate(positions): + for i in range(half): + first, second = (i, i + half) if pairing == "half" else (2 * i, 2 * i + 1) + cos, sin = np.cos(position * theta[i]), np.sin(position * theta[i]) + x_first, x_second = x[(*lead, step, first)], x[(*lead, step, second)] + out[(*lead, step, first)] = x_first * cos - x_second * sin + out[(*lead, step, second)] = x_second * cos + x_first * sin + + return out + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +@pytest.mark.parametrize( + "scaling, scaling_factor", + [("none", 1.0), ("linear", 4.0), ("ntk", 4.0)], + ids=["unscaled", "linear", "ntk"], +) +def test_rope_matches_independent_reference(pairing, scaling, scaling_factor, rng): + x_np = rng.normal(size=(2, 3, 5, 8)).astype(floatX) + positions_np = np.arange(5) + + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + out = rotary_embedding( + x, positions, pairing=pairing, scaling=scaling, scaling_factor=scaling_factor + ) + + result = out.eval({x: x_np, positions: positions_np}) + expected = rope_np( + x_np, positions_np, pairing=pairing, scaling=scaling, scaling_factor=scaling_factor + ) + np.testing.assert_allclose(result, expected, atol=1e-6) + + +def test_rope_pairing_conventions_are_not_interchangeable(rng): + # The pairing is the easiest thing to get silently wrong -- both conventions produce a plausible + # rotation of the right shape and norm. Each must match its own reference and no other. + x_np = rng.normal(size=(4, 6)).astype(floatX) + positions_np = np.arange(1, 5) + + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + half = rotary_embedding(x, positions, pairing="half").eval({x: x_np, positions: positions_np}) + adjacent = rotary_embedding(x, positions, pairing="adjacent").eval( + {x: x_np, positions: positions_np} + ) + + assert not np.allclose(half, adjacent) + np.testing.assert_allclose(half, rope_np(x_np, positions_np, pairing="half"), atol=1e-6) + np.testing.assert_allclose(adjacent, rope_np(x_np, positions_np, pairing="adjacent"), atol=1e-6) + # Explicitly: swapping the convention on the reference side must fail, so a matching test cannot + # pass on a flipped implementation. + assert not np.allclose(half, rope_np(x_np, positions_np, pairing="adjacent")) + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_rope_one_step_at_a_time_equals_whole_sequence(pairing, rng): + # The decode invariant: a single token rotated by its absolute position, with the sequence axis + # length 1, must equal that token's slice of the full-sequence result. Without it, cached decoding + # would silently disagree with prefill. + x_np = rng.normal(size=(2, 3, 6, 4)).astype(floatX) + positions_np = np.arange(6) + + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + whole = rotary_embedding(x, positions, pairing=pairing).eval({x: x_np, positions: positions_np}) + + step_x = pt.tensor("step_x", shape=(2, 3, 1, 4)) + step_out = rotary_embedding(step_x, positions, pairing=pairing) + for step, position in enumerate(positions_np): + one = step_out.eval({step_x: x_np[:, :, step : step + 1], positions: np.array([position])}) + np.testing.assert_allclose(one[:, :, 0], whole[:, :, step], atol=1e-6) + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_rope_scores_depend_only_on_relative_position(pairing, rng): + # The property RoPE exists for: because each rotation is orthogonal, the dot product of a rotated + # query and a rotated key is a function of the position difference alone. A wrong sign or a + # mispaired channel breaks this while leaving norms intact. + q_np = rng.normal(size=(1, 1, 1, 8)).astype(floatX) + k_np = rng.normal(size=(1, 1, 1, 8)).astype(floatX) + + q, k = pt.tensor("q", shape=q_np.shape), pt.tensor("k", shape=k_np.shape) + positions = pt.lvector("positions") + q_out = rotary_embedding(q, positions, pairing=pairing) + k_out = rotary_embedding(k, positions, pairing=pairing) + + def score(query_position, key_position): + rotated_q = q_out.eval({q: q_np, positions: np.array([query_position])}) + rotated_k = k_out.eval({k: k_np, positions: np.array([key_position])}) + return float((rotated_q * rotated_k).sum()) + + np.testing.assert_allclose(score(3, 0), score(10, 7), rtol=1e-6) + np.testing.assert_allclose(score(3, 0), score(103, 100), rtol=1e-6) + assert not np.isclose(score(3, 0), score(4, 0)) + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_rope_is_an_orthogonal_rotation(pairing, rng): + # Rotations preserve length, and position 0 rotates by zero radians. + x_np = rng.normal(size=(2, 4, 6)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + out = rotary_embedding(x, positions, pairing=pairing) + + rotated = out.eval({x: x_np, positions: np.arange(4)}) + np.testing.assert_allclose( + np.linalg.norm(rotated, axis=-1), np.linalg.norm(x_np, axis=-1), rtol=1e-6 + ) + + identity = out.eval({x: x_np, positions: np.zeros(4, dtype="int64")}) + np.testing.assert_allclose(identity, x_np, atol=1e-6) + + +def test_rope_positions_broadcast_shared_and_per_sequence(rng): + # Positions carry only batch axes and the sequence axis; the head axes of x are supplied by the + # layer, so (seq,) and (batch, seq) both line up with (batch, n_head, seq, head_dim). + x_np = rng.normal(size=(2, 3, 5, 8)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + + shared = pt.lvector("shared") + per_sequence = pt.lmatrix("per_sequence") + shared_out = rotary_embedding(x, shared) + batched_out = rotary_embedding(x, per_sequence) + + positions_np = np.arange(5) + shared_result = shared_out.eval({x: x_np, shared: positions_np}) + repeated = batched_out.eval({x: x_np, per_sequence: np.tile(positions_np, (2, 1))}) + np.testing.assert_allclose(shared_result, repeated, atol=1e-6) + + # Independent offsets per batch row: row 0 keeps the shared answer, row 1 must not. + offsets = np.stack([positions_np, positions_np + 100]) + offset_result = batched_out.eval({x: x_np, per_sequence: offsets}) + np.testing.assert_allclose(offset_result[0], shared_result[0], atol=1e-6) + assert not np.allclose(offset_result[1], shared_result[1]) + np.testing.assert_allclose(offset_result[1], rope_np(x_np[1], positions_np + 100), atol=1e-6) + + +@pytest.mark.parametrize("scaling", ["linear", "ntk"]) +def test_rope_scaling_factor_one_is_the_identity(scaling, rng): + x_np = rng.normal(size=(3, 6)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + positions_np = np.arange(3) + + unscaled = rotary_embedding(x, positions).eval({x: x_np, positions: positions_np}) + scaled = rotary_embedding(x, positions, scaling=scaling, scaling_factor=1.0).eval( + {x: x_np, positions: positions_np} + ) + np.testing.assert_allclose(unscaled, scaled, atol=1e-6) + + +def test_linear_scaling_interpolates_positions(rng): + # Position interpolation divides the frequencies, which is the same as dividing the positions: + # rotating position 8 with factor 4 must equal rotating position 2 unscaled. + x_np = rng.normal(size=(1, 6)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + + interpolated = rotary_embedding(x, positions, scaling="linear", scaling_factor=4.0).eval( + {x: x_np, positions: np.array([8])} + ) + plain = rotary_embedding(x, positions).eval({x: x_np, positions: np.array([2])}) + np.testing.assert_allclose(interpolated, plain, atol=1e-6) + + +def test_rope_composes_with_attention_unchanged(rng): + # RoPE acts on q and k before attention, so the attention op needs no knowledge of positions. + # Rotating both must change the scores (it is not a no-op) while leaving the output shape alone. + q_np, k_np, v_np = (rng.normal(size=(2, 3, 5, 4)).astype(floatX) for _ in range(3)) + q = pt.tensor("q", shape=q_np.shape) + k = pt.tensor("k", shape=k_np.shape) + v = pt.tensor("v", shape=v_np.shape) + positions = pt.lvector("positions") + + rope = RotaryEmbedding("rope") + with_rope = scaled_dot_product_attention( + rope(q, positions), rope(k, positions), v, is_causal=True + ) + without_rope = scaled_dot_product_attention(q, k, v, is_causal=True) + + attention_values = {q: q_np, k: k_np, v: v_np} + baseline = without_rope.eval(attention_values) + + rotated = with_rope.eval({**attention_values, positions: np.arange(5)}) + assert rotated.shape == (2, 3, 5, 4) + assert not np.allclose(rotated, baseline) + + # Rotating with all-zero positions is the identity, so attention must see the original q and k. + np.testing.assert_allclose( + with_rope.eval({**attention_values, positions: np.zeros(5, dtype="int64")}), + baseline, + atol=1e-6, + ) + + +def test_rope_layer_carries_its_configuration(rng): + # The layer exists so queries and keys are guaranteed to be rotated with the same frequencies. + x_np = rng.normal(size=(2, 4)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + + rope = RotaryEmbedding("rope", base=500.0, pairing="adjacent") + out = rope(x, positions) + + assert out.name == "rope_output" + np.testing.assert_allclose( + out.eval({x: x_np, positions: np.arange(2)}), + rope_np(x_np, np.arange(2), base=500.0, pairing="adjacent"), + atol=1e-6, + ) + + +def test_rope_prediction_matches_training(rng): + # RoPE is a fixed function of position, identical in train and eval, so like LayerNorm it needs no + # prediction rewrite: rewrite_for_prediction must leave its output untouched. + x_np = rng.normal(size=(2, 6)).astype(floatX) + x = pt.tensor("x", shape=x_np.shape) + positions = pt.lvector("positions") + out = rotary_embedding(x, positions) + + values = {x: x_np, positions: np.arange(2)} + np.testing.assert_allclose( + rewrite_for_prediction(out).eval(values), out.eval(values), rtol=1e-6 + ) + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_rope_passes_gradients_to_its_input(pairing, rng): + # The interleaved branch reassembles its output with strided writes, whose gradient has to route + # every channel back; a dropped slice would show up as a zero column here. + x = pt.tensor("x", shape=(2, 4)) + positions = pt.lvector("positions") + out = rotary_embedding(x, positions, pairing=pairing) + + x_np = rng.normal(size=(2, 4)).astype(floatX) + grad = pt.grad(out.sum(), x).eval({x: x_np, positions: np.arange(2)}) + assert np.all(np.isfinite(grad)) + assert not np.any(np.all(np.isclose(grad, 0.0), axis=0)) + + +def test_unknown_head_dimension_raises(): + with pytest.raises(ValueError, match="statically known last dimension"): + rotary_embedding(pt.matrix("x"), pt.lvector("positions")) + + +def test_odd_head_dimension_raises(): + with pytest.raises(ValueError, match="even head dimension"): + rotary_embedding(pt.tensor("x", shape=(4, 7)), pt.lvector("positions")) + + +def test_integer_input_raises(): + with pytest.raises(ValueError, match="floating-point tensor"): + rotary_embedding(pt.tensor("x", shape=(4, 8), dtype="int64"), pt.lvector("positions")) + + +def test_positions_wider_than_input_raises(): + with pytest.raises(ValueError, match="more than the 1 non-feature dimensions"): + rotary_embedding(pt.tensor("x", shape=(4, 8)), pt.lmatrix("positions")) + + +@pytest.mark.parametrize( + "kwargs, match", + [ + ({"pairing": "interleaved"}, "pairing must be one of"), + ({"scaling": "yarn"}, "scaling must be one of"), + ({"scaling": "linear", "scaling_factor": 0.0}, "scaling_factor must be positive"), + ], + ids=["bad_pairing", "bad_scaling", "bad_factor"], +) +def test_invalid_options_raise(kwargs, match): + with pytest.raises(ValueError, match=match): + rotary_embedding(pt.tensor("x", shape=(4, 8)), pt.lvector("positions"), **kwargs) + with pytest.raises(ValueError, match=match): + RotaryEmbedding("rope", **kwargs) + + +def test_ntk_scaling_needs_more_than_two_channels(): + with pytest.raises(ValueError, match="undefined for head_dim <= 2"): + rotary_embedding( + pt.tensor("x", shape=(4, 2)), pt.lvector("positions"), scaling="ntk", scaling_factor=2.0 + ) diff --git a/tests/test_serialize.py b/tests/test_serialize.py index 2d6334c..22e0cde 100644 --- a/tests/test_serialize.py +++ b/tests/test_serialize.py @@ -27,10 +27,12 @@ Embedding, LayerNorm, Linear, + RMSNorm, Sequential, Squeeze, ) from pytensor_ml.layers.attention import scaled_dot_product_attention +from pytensor_ml.layers.positional import rotary_embedding from pytensor_ml.pytensorf import collect_shared_variables, collect_trainable_params from pytensor_ml.serialize.base import _TYPE_FROM_JSON, _TYPE_TO_JSON @@ -111,6 +113,30 @@ def test_layernorm_roundtrips(affine): assert_outputs_roundtrip([X], output, [np.random.default_rng(1).normal(size=(8, 4))]) +@pytest.mark.parametrize("affine", [True, False], ids=["affine", "no_affine"]) +def test_rmsnorm_roundtrips(affine): + X, output = initialized_network(Linear("fc", 4, 6), RMSNorm("rms", n_in=6, affine=affine)) + assert_outputs_roundtrip([X], output, [np.random.default_rng(1).normal(size=(8, 4))]) + + +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +@pytest.mark.parametrize( + "scaling, scaling_factor", + [("none", 1.0), ("linear", 4.0), ("ntk", 4.0)], + ids=["unscaled", "linear", "ntk"], +) +def test_rotary_embedding_roundtrips(pairing, scaling, scaling_factor): + # Every option is a prop, so the rebuilt op must reproduce the same frequencies; a dropped prop + # would silently fall back to the default and still evaluate. + x = pt.tensor("x", shape=(2, 3, 5, 8)) + positions = pt.lvector("positions") + output = rotary_embedding( + x, positions, base=500.0, pairing=pairing, scaling=scaling, scaling_factor=scaling_factor + ) + values = [np.random.default_rng(0).normal(size=(2, 3, 5, 8)), np.arange(5)] + assert_outputs_roundtrip([x, positions], output, values) + + def test_squeeze_roundtrips(): X = pt.matrix("X") assert_outputs_roundtrip( From 8aeb005ab74cf444c7423bcb02a775a9e474dc2a Mon Sep 17 00:00:00 2001 From: Carlos Trujillo Date: Sat, 8 Aug 2026 18:19:44 +0300 Subject: [PATCH 2/2] Address review: simplify, fix false JAX claim, correct dtype handling Revert the unrelated churn: `activations.py` is untouched again and the `constant_like` move is gone. RMSNorm pins its epsilon with the core one-liner `pt.constant(self.epsilon, dtype=X.type.dtype)` instead. `base.py` now only gains `PositionalLayer`, a sibling ABC for layers taking positions. It is a sibling and not a subclass of `Layer` because narrowing `Layer.__call__` from one tensor to two is not a valid override -- verified that subclassing still errors under mypy, so a subclass would only move the `type: ignore` up a level. The ignore is now gone entirely. Simplifications: `_rms_normalize` inlined into `RMSNormLayer` (single caller), `_scale_parameter` deleted and `_affine_parameters` restored to its symmetric form, dead comments removed. RoPE no longer requires a static `head_dim`. The frequency ladder is built with `pt.arange` and folds to a single constant when the size is known, so the static requirement bought nothing; a symbolic feature axis now works and is tested. Everything is computed in the input's dtype via `x.type.dtype`, so no float64 leaks into single-precision graphs. Returned variables are named. The interleaved pairing uses `x[..., 0::2].set(...)`. Retract a false claim: `SplitDims`/`JoinDims` do lower on JAX -- they are rewritten to `Reshape` during canonicalization. The real defect was in `tests/dispatch/jax/test_basic.py`, whose `compare_jax_and_py` built its own `RewriteDatabaseQuery(include=["jax"])` that omits canonicalize, so any op which only reaches the backend after being rewritten raised "No JAX conversion" there while compiling fine in real use. It now uses `JAX._optimizer`. `MultiheadAttention` is and always was JAX-compilable. Docs: American English, References sections with the RoFormer, positional interpolation and RMSNorm papers, pinned permalinks to the HuggingFace and torchtune implementations for the two pairing conventions, parameter docs cut down, `See Also` removed. --- pytensor_ml/activations.py | 32 +++- pytensor_ml/base.py | 32 ++-- pytensor_ml/layers/norm.py | 52 ++----- pytensor_ml/layers/positional.py | 245 ++++++++++++++----------------- tests/dispatch/jax/test_basic.py | 7 +- tests/test_layers.py | 2 +- tests/test_positional.py | 24 ++- 7 files changed, 183 insertions(+), 211 deletions(-) diff --git a/pytensor_ml/activations.py b/pytensor_ml/activations.py index 186ad52..f1cf253 100644 --- a/pytensor_ml/activations.py +++ b/pytensor_ml/activations.py @@ -1,7 +1,23 @@ import numpy as np import pytensor.tensor as pt -from pytensor_ml.base import Layer, constant_like +from pytensor import config + +from pytensor_ml.base import Layer + + +def _constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable: + """ + Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype. + + PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand + depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x`` + promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence. + """ + dtype = np.dtype(x.dtype) + # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision. + dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX) + return pt.constant(np.asarray(value, dtype=dtype)) class Activation(Layer): ... @@ -46,7 +62,7 @@ def __init__(self, negative_slope: float = 0.01): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - out = pt.switch(x > 0, x, constant_like(self.negative_slope, x) * x) + out = pt.switch(x > 0, x, _constant_like(self.negative_slope, x) * x) out.name = "LeakyReLU" return out @@ -133,14 +149,14 @@ def __init__(self, approximate: bool = True): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - half = constant_like(0.5, x) - one = constant_like(1.0, x) + half = _constant_like(0.5, x) + one = _constant_like(1.0, x) if self.approximate: - sqrt_2_over_pi = constant_like(np.sqrt(2.0 / np.pi), x) - cubic_coef = constant_like(0.044715, x) + sqrt_2_over_pi = _constant_like(np.sqrt(2.0 / np.pi), x) + cubic_coef = _constant_like(0.044715, x) out = half * x * (one + pt.tanh(sqrt_2_over_pi * (x + cubic_coef * x**3))) else: - sqrt2 = constant_like(np.sqrt(2.0), x) + sqrt2 = _constant_like(np.sqrt(2.0), x) out = half * x * (one + pt.erf(x / sqrt2)) out.name = "GELU" return out @@ -166,7 +182,7 @@ def __init__(self, beta: float = 1.0): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: x = pt.as_tensor(x) - out = x * pt.sigmoid(constant_like(self.beta, x) * x) + out = x * pt.sigmoid(_constant_like(self.beta, x) * x) out.name = "Swish" return out diff --git a/pytensor_ml/base.py b/pytensor_ml/base.py index 42a6864..6268c9d 100644 --- a/pytensor_ml/base.py +++ b/pytensor_ml/base.py @@ -1,31 +1,12 @@ from abc import ABC, abstractmethod from typing import Protocol, runtime_checkable -import numpy as np import pytensor.tensor as pt -from pytensor import config from pytensor.compile.builders import SymbolicOp from pytensor.tensor.variable import TensorVariable -def constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable: - """ - Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype. - - PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand - depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x`` - promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence. - - Lives here, beside the layer base classes, because both ``pytensor_ml.activations`` and the layer - modules need it and this module is already the shared root that neither can cycle through. - """ - dtype = np.dtype(x.dtype) - # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision. - dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX) - return pt.constant(np.asarray(value, dtype=dtype)) - - class Layer(ABC): """Base class for the objects that build layer graphs. Defined here, not in ``pytensor_ml.layers``, so that ``pytensor_ml.activations`` can subclass it without a circular import.""" @@ -34,6 +15,17 @@ class Layer(ABC): def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: ... +class PositionalLayer(ABC): + """Base class for layers that consume token positions alongside their input. + + A sibling of :class:`Layer` rather than a subclass: narrowing ``Layer.__call__`` from one tensor to + two is not a valid override, so a subclass would need a type-checker escape at every level. + """ + + @abstractmethod + def __call__(self, x: pt.TensorLike, position_ids: pt.TensorLike) -> pt.TensorVariable: ... + + class LayerOp(SymbolicOp): """Base class for the library's neural-network ops. @@ -72,4 +64,4 @@ def update_map(self) -> dict[int, int]: """Map each output index to the index of the input that output updates.""" -__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp", "constant_like"] +__all__ = ["Layer", "LayerOp", "PositionalLayer", "StatefulOp", "UnaryLayerOp"] diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py index 8c688ce..bff0b7e 100644 --- a/pytensor_ml/layers/norm.py +++ b/pytensor_ml/layers/norm.py @@ -3,7 +3,7 @@ from pytensor import config -from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp, constant_like +from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp from pytensor_ml.params import NonTrainableParameter, TrainableParameter, non_trainable, trainable @@ -31,21 +31,6 @@ def _standardize(X, epsilon, axis, keepdims=False): return (X - mu) / pt.sqrt(sigma_sq + epsilon), mu, sigma_sq -def _rms_normalize(X, epsilon): - """ - Scale ``X`` to unit root mean square over its last axis, without centering it. - - Notes - ----- - RMSNorm divides by the root mean square rather than the standard deviation, so unlike - :func:`_standardize` it leaves the mean of ``X`` intact. That omission is the entire difference - between the two, and pretrained weights depend on it -- do not "simplify" this into a - :func:`_standardize` call. - """ - mean_square = pt.mean(pt.square(X), axis=-1, keepdims=True) - return X / pt.sqrt(mean_square + constant_like(epsilon, X)) - - def _affine_input_count(affine: bool) -> int: """Number of inputs the learned affine transform contributes, which every norm op places directly after ``X``. Both the graph builders and :meth:`BatchNormLayer.update_map` index around it.""" @@ -90,17 +75,12 @@ def _resolve_n_in(name: str, n_in: int | None, X: pt.TensorVariable | None) -> i return inferred -def _scale_parameter(name: str, n_in: int) -> TrainableParameter: - """Build the learned scale. RMSNorm has only this one; the shift-and-scale norms pair it with a - ``loc``, so the ``_scale`` suffix is fixed here and nowhere else.""" - return trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale") - - def _affine_parameters(name: str, n_in: int) -> tuple[TrainableParameter, TrainableParameter]: """Build the learned shift and scale. Returns them in the ``(loc, scale)`` order that every norm op unpacks its inputs in, so the two cannot drift apart.""" loc = trainable(np.zeros(n_in, dtype=config.floatX), f"{name}_loc") - return loc, _scale_parameter(name, n_in) + scale = trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale") + return loc, scale class BatchNormLayer(LayerOp): @@ -361,12 +341,12 @@ class RMSNormLayer(UnaryLayerOp): __props__ = ("n_in", "epsilon", "affine") def build_inner_graph(self, X, *rest): - X_normalized = _rms_normalize(X, self.epsilon) + epsilon = pt.constant(self.epsilon, dtype=X.type.dtype) + mean_square = pt.mean(pt.square(X), axis=-1, keepdims=True) + X_normalized = X / pt.sqrt(mean_square + epsilon) if not self.affine: return [X_normalized] - # Scale-only, so the affine transform contributes one input rather than the ``(loc, scale)`` - # pair the other norm ops unpack; the single-element unpack asserts that arity. (scale,) = rest return [X_normalized * scale] @@ -382,10 +362,8 @@ class RMSNorm(Layer): y = \frac{x}{\sqrt{\frac{1}{n} \sum_i x_i^2 + \epsilon}} \cdot \gamma. - Unlike :class:`LayerNorm` there is no mean subtraction and no learned shift. That is not a - simplification of this implementation but the definition: :class:`torch.nn.RMSNorm`, - ``flax.linen.RMSNorm``/``flax.nnx.RMSNorm`` and ``tinygrad.nn.RMSNorm`` all expose a weight and - no bias. It is the normalization used by the Llama, Gemma and Qwen decoder families. + Unlike :class:`LayerNorm` there is no mean subtraction and no learned shift; the scale is the only + parameter. This is the normalization used by the Llama, Gemma and Qwen decoder families. Parameters ---------- @@ -395,12 +373,14 @@ class RMSNorm(Layer): Size of the normalized feature axis. Inferred from the input's last dimension on the first call when omitted. epsilon : float, optional - Constant :math:`\epsilon` added to the mean square for numerical stability. Default is 1e-6, - matching ``flax``'s and ``tinygrad``'s RMSNorm. This deliberately differs from - :class:`LayerNorm`'s 1e-5, which follows :class:`torch.nn.LayerNorm`; pretrained weights are - published against their own framework's value. + Constant :math:`\epsilon` added to the mean square for numerical stability. Default is 1e-6. affine : bool, optional - Apply the learned scale :math:`\gamma`. There is no shift to disable. Default is True. + Apply the learned scale :math:`\gamma`. Default is True. + + References + ---------- + .. [1] Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. + arXiv:1910.07467. https://arxiv.org/abs/1910.07467. """ def __init__( @@ -429,7 +409,7 @@ def _initialize_params(self, X: pt.TensorVariable | None): return if self.affine: - self.scale = _scale_parameter(self.name, n_in) + self.scale = trainable(np.ones(n_in, dtype=config.floatX), f"{self.name}_scale") self.initialized = True diff --git a/pytensor_ml/layers/positional.py b/pytensor_ml/layers/positional.py index 5110075..3789ea0 100644 --- a/pytensor_ml/layers/positional.py +++ b/pytensor_ml/layers/positional.py @@ -1,11 +1,11 @@ from typing import Literal, get_args -import numpy as np import pytensor.tensor as pt +from pytensor.tensor.type import float_dtypes from pytensor.tensor.variable import TensorVariable -from pytensor_ml.base import Layer, UnaryLayerOp +from pytensor_ml.base import PositionalLayer, UnaryLayerOp Pairing = Literal["half", "adjacent"] Scaling = Literal["none", "linear", "ntk"] @@ -15,8 +15,8 @@ def _validate_options(pairing: str, scaling: str, scaling_factor: float) -> None: - """Reject unknown options where they are given, rather than letting a typo silently select the - default behaviour inside the inner graph.""" + """Reject unknown options where they are given, so a typo cannot silently select the default + behavior inside the inner graph.""" if pairing not in _PAIRINGS: raise ValueError(f"pairing must be one of {_PAIRINGS}, got {pairing!r}") if scaling not in _SCALINGS: @@ -25,50 +25,40 @@ def _validate_options(pairing: str, scaling: str, scaling_factor: float) -> None raise ValueError(f"scaling_factor must be positive, got {scaling_factor}") -def _resolve_head_dim(x: TensorVariable) -> int: - """The rotated feature size, which must be static and even. +def _head_dim(x: TensorVariable) -> int | TensorVariable: + """Size of the rotated feature axis, as a Python int when ``x`` declares one. - Static because the frequencies are materialized as a constant at graph-build time, and even - because the rotation acts on two-dimensional subspaces. + A static size lets the frequency ladder fold to a constant and keeps the output's static shape; + a symbolic one works too, so it is a preference rather than a requirement. """ - if not np.issubdtype(np.dtype(x.dtype), np.floating): - raise ValueError( - f"RotaryEmbedding rotates a floating-point tensor, but got dtype {x.dtype}. Cast the " - f"input first." - ) + if x.type.dtype not in float_dtypes: + raise ValueError(f"RotaryEmbedding needs a floating-point input, got dtype {x.type.dtype}.") - head_dim = x.type.shape[-1] - if head_dim is None: + static_size = x.type.shape[-1] + if static_size is None: + return x.shape[-1] + if static_size % 2: raise ValueError( - "RotaryEmbedding needs a statically known last dimension to build its frequencies. " - "Give the input a static feature dimension, e.g. pt.tensor(shape=(None, n_head, None, " - "head_dim))." + f"RotaryEmbedding rotates channel pairs, so head_dim must be even, got {static_size}." ) - if head_dim % 2: - raise ValueError(f"RotaryEmbedding needs an even head dimension, got {head_dim}.") - return head_dim + return static_size -def _align_to_head_axes( +def _add_head_axes( angles: TensorVariable, x: TensorVariable, position_ids: TensorVariable ) -> TensorVariable: - """Insert ``x``'s head axes into the angle tensor so the sequence axes line up. - - ``position_ids`` describes tokens: its last axis is the sequence and any leading axes are batch - axes. ``x`` carries its head axes *between* those two, so the angles need matching placeholder - axes. How many follows from the ranks, which is why the caller does not pass an unsqueeze position: - ``(seq,)`` broadcasts over everything, and ``(batch, seq)`` lines up against - ``(batch, n_head, seq, head_dim)`` without the caller reshaping. Deriving it also keeps the - inserted axes statically 1, which is what pytensor requires before it will broadcast them -- a - hand-built ``(batch, 1, seq)`` input whose middle axis is merely unknown is rejected at runtime. + """Unsqueeze ``angles`` so its sequence axis lines up with ``x``'s. + + ``position_ids`` indexes tokens, so its last axis is the sequence and any leading axes are batch + axes. ``x`` carries head axes between the two, and the number of them follows from the two ranks -- + which is why the caller does not pass an axis to unsqueeze. """ n_head_axes = (x.type.ndim - 1) - position_ids.type.ndim if n_head_axes < 0: raise ValueError( f"position_ids has {position_ids.type.ndim} dimensions, more than the " - f"{x.type.ndim - 1} non-feature dimensions of x. Its last axis is the sequence and any " - f"leading axes are batch axes, so it can never be wider than x minus the feature axis." + f"{x.type.ndim - 1} non-feature dimensions of x." ) if n_head_axes == 0: return angles @@ -77,66 +67,55 @@ def _align_to_head_axes( def _inverse_frequencies( - head_dim: int, base: float, scaling: str, scaling_factor: float -) -> np.ndarray: + head_dim: int | TensorVariable, dtype: str, base: float, scaling: str, scaling_factor: float +) -> TensorVariable: r""" - Angular frequencies :math:`\theta_i = \mathrm{base}^{-2i/d}` for each rotated pair. - - Computed in numpy at graph-build time: ``base``, ``head_dim`` and the scaling are all known then, - so the frequencies enter the graph as one constant instead of a subgraph of ``arange`` and - ``power``. - - Parameters - ---------- - head_dim : int - Size of the rotated feature axis. There are ``head_dim // 2`` frequencies. - base : float - Geometric base of the frequency ladder, :math:`\theta` in the RoFormer paper. - scaling : str - ``"none"``, ``"linear"`` for position interpolation, or ``"ntk"`` for NTK-aware scaling. - scaling_factor : float - Context-extension factor for the scaled variants; ignored when ``scaling="none"``. + Angular frequencies :math:`\theta_i = \mathrm{base}^{-2i/d}`, one per rotated pair. Returns ------- - ndarray - Frequencies of shape ``(head_dim // 2,)``, in float64. + inverse_frequencies : TensorVariable + Shape ``(head_dim // 2,)``, dtype ``dtype``. Folds to a constant when ``head_dim`` is static. """ if scaling == "ntk": - if head_dim <= 2: + if isinstance(head_dim, int) and head_dim <= 2: raise ValueError( - f"NTK-aware scaling rescales the base by scaling_factor ** (d / (d - 2)), which is " + f"NTK scaling rescales the base by scaling_factor ** (d / (d - 2)), which is " f"undefined for head_dim <= 2; got {head_dim}." ) - # Static NTK-aware scaling: stretch the frequency ladder itself rather than the positions, so - # the highest frequencies are left almost untouched. This is the exponent HuggingFace's - # `_compute_dynamic_ntk_parameters` uses; the dynamic form, which recomputes the base from the - # running sequence length, needs a runtime length and is deliberately not offered here. + # Stretches the frequency ladder itself rather than the positions, leaving the highest + # frequencies nearly untouched. Only the static form is offered: the dynamic variant rescales + # the base from the running sequence length, which is not known at graph-build time. base = base * scaling_factor ** (head_dim / (head_dim - 2)) - inverse_frequencies = 1.0 / base ** (np.arange(0, head_dim, 2, dtype="float64") / head_dim) + exponent = pt.arange(0, head_dim, 2, dtype=dtype) / head_dim + inverse_frequencies = base**-exponent if scaling == "linear": - # Position interpolation: dividing the frequencies is algebraically the same as dividing the - # positions, because the angle is bilinear in the two. + # Dividing the frequencies is algebraically the same as dividing the positions, since the + # angle is bilinear in the two. inverse_frequencies = inverse_frequencies / scaling_factor + inverse_frequencies = inverse_frequencies.astype(dtype) + inverse_frequencies.name = "inverse_frequencies" + return inverse_frequencies def _split_pairs( - x: TensorVariable, pairing: str, half: int + x: TensorVariable, pairing: str, head_dim: int | TensorVariable ) -> tuple[TensorVariable, TensorVariable]: - """Split the feature axis into the two halves of every rotated pair. + """Split the feature axis into the two members of every rotated pair. - The two conventions differ only in which channels are paired, but they are not interchangeable: - weights trained under one produce nonsense under the other. + The conventions are a permutation of the feature axis apart and are not interchangeable: weights + trained under one produce nonsense under the other. - ``"half"`` pairs channel ``i`` with ``i + d/2``, matching HuggingFace's ``rotate_half`` (Llama, - Qwen, Mistral, Gemma), GPT-NeoX and flax's ``RoPE``. ``"adjacent"`` pairs ``2i`` with ``2i + 1``, - matching the original RoFormer formulation, GPT-J and ``torchtune``. + ``"half"`` pairs channel ``i`` with ``i + d/2``, as HuggingFace's ``rotate_half`` does. + ``"adjacent"`` pairs ``2i`` with ``2i + 1``, as the original RoFormer paper and torchtune do. Both + are cited from :func:`rotary_embedding`. """ if pairing == "half": + half = head_dim // 2 return x[..., :half], x[..., half:] return x[..., 0::2], x[..., 1::2] @@ -145,21 +124,13 @@ def _split_pairs( def _join_pairs( x: TensorVariable, first: TensorVariable, second: TensorVariable, pairing: str ) -> TensorVariable: - """Reassemble the feature axis, inverting :func:`_split_pairs` for the same ``pairing``. - - Both branches stay on ops every backend lowers, and both preserve the static feature size that - lets the backends vectorize the rotation. Notably ``split_dims``/``join_dims`` would read more - naturally for the interleaved case but have no JAX or MLX conversion, so a graph using them - compiles only in the default backend. - """ + """Reassemble the feature axis, inverting :func:`_split_pairs` for the same ``pairing``.""" if pairing == "half": return pt.concatenate([first, second], axis=-1) - # Strided writes into a copy of x rather than an interleaving reshape. Every channel is - # overwritten -- the even ones by `first`, the odd ones by `second` -- so no value of x survives, - # and both were read before either write. - rotated = pt.set_subtensor(x[..., 0::2], first) - return pt.set_subtensor(rotated[..., 1::2], second) + # Every channel is overwritten -- the even ones by `first`, the odd ones by `second` -- and both + # were read before either write, so nothing of x survives into the result. + return x[..., 0::2].set(first)[..., 1::2].set(second) class RotaryEmbeddingLayer(UnaryLayerOp): @@ -167,23 +138,21 @@ class RotaryEmbeddingLayer(UnaryLayerOp): def build_inner_graph(self, x, position_ids): _validate_options(self.pairing, self.scaling, self.scaling_factor) - head_dim = _resolve_head_dim(x) + head_dim = _head_dim(x) + dtype = x.type.dtype inverse_frequencies = _inverse_frequencies( - head_dim, self.base, self.scaling, self.scaling_factor - ) - # Angles, and therefore cos/sin, are computed in x's own dtype: floatX is float32 or float64 - # here, so this already is the "compute in fp32" path that the fp16/bf16 frameworks upcast to. - angles = position_ids[..., None].astype(x.dtype) * pt.constant( - inverse_frequencies.astype(x.dtype), name="inverse_frequencies" + head_dim, dtype, self.base, self.scaling, self.scaling_factor ) - angles = _align_to_head_axes(angles, x, position_ids) + angles = position_ids[..., None].astype(dtype) * inverse_frequencies + angles = _add_head_axes(angles, x, position_ids) cos, sin = pt.cos(angles), pt.sin(angles) - first, second = _split_pairs(x, self.pairing, head_dim // 2) + first, second = _split_pairs(x, self.pairing, head_dim) rotated = _join_pairs( x, first * cos - second * sin, second * cos + first * sin, self.pairing ) + rotated.name = "rotary_embedding" return [rotated] @@ -211,88 +180,92 @@ def rotary_embedding( \begin{pmatrix} x_a \\ x_b \end{pmatrix}, \qquad \theta_i = \mathrm{base}^{-2i/d}, - where :math:`m` is the position and :math:`(a, b)` is the ``i``-th channel pair. Because a - rotation is orthogonal, the dot product of a rotated query and a rotated key depends only on the - *difference* of their positions -- that relative-position property is the point of the scheme, and - it is what makes RoPE usable for incremental decoding: a token rotated by its absolute position - keeps the same relationship to every earlier token no matter when it was computed. + where :math:`m` is the position and :math:`(a, b)` is the :math:`i`-th channel pair [1]_. The + rotation is orthogonal, so the dot product of a rotated query and a rotated key depends only on + the difference of their positions. That is what makes it usable for incremental decoding: a token + rotated by its absolute position keeps the same relationship to every earlier token however late + it is computed. - Apply this to queries and keys before :func:`~pytensor_ml.layers.attention.scaled_dot_product_attention`, - never to values. + Apply to queries and keys before + :func:`~pytensor_ml.layers.attention.scaled_dot_product_attention`, never to values. Parameters ---------- x : TensorLike Tensor whose last axis is rotated, typically queries or keys of shape - ``(..., n_head, seq, head_dim)``. ``head_dim`` must be static and even. + ``(..., n_head, seq, head_dim)``. ``head_dim`` must be even. position_ids : TensorLike - Position of each token, shape ``(..., seq)``: the last axis is the sequence and any leading - axes are batch axes. ``x``'s head axes are inserted for you, so both ``(seq,)`` -- shared by - the whole batch -- and ``(batch, seq)`` work against ``(batch, n_head, seq, head_dim)``. - Positions are explicit rather than an implied ``0..seq-1`` so that the same graph serves a full - sequence and a single decode step appended to a cached prefix. + Integer position of each token, shape ``(..., seq)``. ``x``'s head axes are unsqueezed in, so + ``(seq,)`` and ``(batch, seq)`` both align with ``(batch, n_head, seq, head_dim)``. Explicit + rather than an implied ``0..seq-1``, so one graph serves a full sequence and a single decode + step appended to a cached prefix. base : float, optional - Geometric base of the frequency ladder (:math:`\theta` in some formulations). Default is - 10000.0, the value used by RoFormer, Llama, GPT-NeoX and torchtune. + Geometric base of the frequency ladder. Default 10000.0. pairing : str, optional - Which channels form each rotated pair. ``"half"`` (default) pairs ``i`` with ``i + d/2``, as - HuggingFace's ``rotate_half``, GPT-NeoX and flax's ``RoPE`` do; ``"adjacent"`` pairs ``2i`` - with ``2i + 1``, as the original RoFormer, GPT-J and torchtune do. The two are a permutation of - the feature axis apart, so this must match whatever the weights were trained with. + Which channels form each rotated pair, ``"half"`` (default) or ``"adjacent"``. Must match the + convention the weights were trained with; see :func:`_split_pairs`. scaling : str, optional - Context-extension scheme. ``"none"`` (default) for plain RoPE, ``"linear"`` for position - interpolation (Chen et al. 2023), or ``"ntk"`` for static NTK-aware scaling, which rescales the - base by ``scaling_factor ** (head_dim / (head_dim - 2))``. + Context-extension scheme: ``"none"`` (default), ``"linear"`` for position interpolation [2]_, + or ``"ntk"`` for static NTK-aware scaling. scaling_factor : float, optional - Extension factor for the scaled variants, ignored when ``scaling="none"``. Default is 1.0, - which is the identity for both variants. + Extension factor for the scaled variants, ignored when ``scaling="none"``. Default 1.0, the + identity for both. Returns ------- - TensorVariable + rotated : TensorVariable ``x`` with its last axis rotated, same shape and dtype. + + References + ---------- + .. [1] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced + Transformer with Rotary Position Embedding. arXiv:2104.09864. https://arxiv.org/abs/2104.09864. + .. [2] Chen, S., Wong, S., Chen, L., & Tian, Y. (2023). Extending Context Window of Large Language + Models via Positional Interpolation. arXiv:2306.15595. https://arxiv.org/abs/2306.15595. + .. [3] https://github.com/huggingface/transformers/blob/v4.57.6/src/transformers/models/llama/modeling_llama.py#L109-L113 + .. [4] https://github.com/meta-pytorch/torchtune/blob/v0.6.1/torchtune/modules/position_embeddings.py#L99-L113 """ _validate_options(pairing, scaling, scaling_factor) x = pt.as_tensor(x) position_ids = pt.as_tensor(position_ids) - result = RotaryEmbeddingLayer( + rotated = RotaryEmbeddingLayer( name="RotaryEmbedding", base=base, pairing=pairing, scaling=scaling, scaling_factor=scaling_factor, )(x, position_ids) - result.name = "rotary_embedding_output" + rotated.name = "rotary_embedding_output" - return result + return rotated -class RotaryEmbedding(Layer): +class RotaryEmbedding(PositionalLayer): r""" - Rotary position embeddings as a reusable, configured layer. + Rotary position embeddings as a configured layer. - Holds the frequency configuration so that queries and keys are rotated identically -- they must be, - since attention compares them -- and calls :func:`rotary_embedding` to build the graph. The layer - has no parameters of its own. + Holds the frequency configuration so queries and keys are rotated identically -- they must be, + since attention compares them. The layer has no parameters of its own. Parameters ---------- name : str or None Name prefix for the layer's output. Defaults to "RotaryEmbedding" when None. base : float, optional - Geometric base of the frequency ladder. Default is 10000.0. + Geometric base of the frequency ladder. Default 10000.0. pairing : str, optional - ``"half"`` or ``"adjacent"``. Default is ``"half"``. + ``"half"`` (default) or ``"adjacent"``. scaling : str, optional - ``"none"``, ``"linear"``, or ``"ntk"``. Default is ``"none"``. + ``"none"`` (default), ``"linear"``, or ``"ntk"``. scaling_factor : float, optional - Extension factor for the scaled variants. Default is 1.0. + Extension factor for the scaled variants. Default 1.0. - See Also - -------- - rotary_embedding : The functional form, documenting the conventions in full. + References + ---------- + .. [1] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced + Transformer with Rotary Position Embedding. arXiv:2104.09864. https://arxiv.org/abs/2104.09864. """ def __init__( @@ -311,12 +284,8 @@ def __init__( self.scaling = scaling self.scaling_factor = scaling_factor - # Positions are a required second input, which the one-tensor Layer.__call__ signature does not - # describe. Widening the base class would loosen it for every layer that really is unary. - def __call__( # type: ignore[override] - self, x: pt.TensorLike, position_ids: pt.TensorLike - ) -> TensorVariable: - out = rotary_embedding( + def __call__(self, x: pt.TensorLike, position_ids: pt.TensorLike) -> TensorVariable: + rotated = rotary_embedding( x, position_ids, base=self.base, @@ -324,9 +293,9 @@ def __call__( # type: ignore[override] scaling=self.scaling, scaling_factor=self.scaling_factor, ) - out.name = f"{self.name}_output" + rotated.name = f"{self.name}_output" - return out + return rotated __all__ = [ diff --git a/tests/dispatch/jax/test_basic.py b/tests/dispatch/jax/test_basic.py index 851aff5..c496917 100644 --- a/tests/dispatch/jax/test_basic.py +++ b/tests/dispatch/jax/test_basic.py @@ -7,13 +7,14 @@ from pytensor.compile.mode import JAX, Mode from pytensor.graph.basic import Variable -from pytensor.graph.rewriting.db import RewriteDatabaseQuery from pytensor.link.jax.linker import JAXLinker jax = pytest.importorskip("jax") -optimizer = RewriteDatabaseQuery(include=["jax"], exclude=JAX._optimizer.exclude) -jax_mode = Mode(linker=JAXLinker(), optimizer=optimizer) +# The JAX mode's own optimizer, not a hand-built ``include=["jax"]`` query. A narrower query skips +# canonicalize, so ops that reach the backend only after being rewritten -- SplitDims and JoinDims +# become Reshape, for instance -- raise "No JAX conversion" here while compiling fine in real use. +jax_mode = Mode(linker=JAXLinker(), optimizer=JAX._optimizer) py_mode = Mode(linker="py", optimizer=None) diff --git a/tests/test_layers.py b/tests/test_layers.py index 90e6907..2fb0a8d 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -265,7 +265,7 @@ def test_rms_norm_no_affine_gives_unit_root_mean_square(rng): np.testing.assert_allclose(np.sqrt(np.square(res).mean(axis=-1)), 1.0, rtol=1e-3) -def test_rms_norm_does_not_centre_its_input(rng): +def test_rms_norm_does_not_center_its_input(rng): # The one thing that separates RMSNorm from LayerNorm: the mean survives. On input with a large # offset, LayerNorm removes it and RMSNorm does not, so the two must disagree. X = pt.tensor("X", shape=(None, 8)) diff --git a/tests/test_positional.py b/tests/test_positional.py index 2876043..89afd8f 100644 --- a/tests/test_positional.py +++ b/tests/test_positional.py @@ -286,18 +286,32 @@ def test_rope_passes_gradients_to_its_input(pairing, rng): assert not np.any(np.all(np.isclose(grad, 0.0), axis=0)) -def test_unknown_head_dimension_raises(): - with pytest.raises(ValueError, match="statically known last dimension"): - rotary_embedding(pt.matrix("x"), pt.lvector("positions")) +@pytest.mark.parametrize("pairing", ["half", "adjacent"]) +def test_unknown_head_dimension_is_supported(pairing, rng): + # The frequency ladder is built symbolically, so a feature axis that is only known at runtime + # works; it just cannot carry a static output shape or fold to a constant. + x_np = rng.normal(size=(3, 8)).astype(floatX) + positions_np = np.arange(3) + + x = pt.tensor("x", shape=(3, None)) + positions = pt.lvector("positions") + out = rotary_embedding(x, positions, pairing=pairing) + + assert out.type.shape == (3, None) + np.testing.assert_allclose( + out.eval({x: x_np, positions: positions_np}), + rope_np(x_np, positions_np, pairing=pairing), + atol=1e-6, + ) def test_odd_head_dimension_raises(): - with pytest.raises(ValueError, match="even head dimension"): + with pytest.raises(ValueError, match="head_dim must be even"): rotary_embedding(pt.tensor("x", shape=(4, 7)), pt.lvector("positions")) def test_integer_input_raises(): - with pytest.raises(ValueError, match="floating-point tensor"): + with pytest.raises(ValueError, match="floating-point input"): rotary_embedding(pt.tensor("x", shape=(4, 8), dtype="int64"), pt.lvector("positions"))