Add RMSNorm and rotary position embeddings by cetagostini · Pull Request #55 · pymc-devs/pytensor_ml

cetagostini · 2026-08-07T20:29:08Z

Two additive decoder primitives from the in-scope set in the #53 triage, nothing else. Both are UnaryLayerOps 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, no change to any existing suite.

RMSNorm

RMSNorm / RMSNormLayer, beside the existing LayerNorm / LayerNormLayer. Scale-only affine: no shift, no mean subtraction. epsilon=1e-6. Reference: Zhang & Sennrich, arXiv:1910.07467.

RoPE

New pytensor_ml/layers/positional.py: rotary_embedding (functional, mirroring scaled_dot_product_attention), RotaryEmbedding (holds the configuration so q and k are rotated identically) and RotaryEmbeddingLayer (marker op).

Positions are an explicit input, not an implied 0..seq-1, so one graph serves a full sequence and a single decode step appended to a prefix. Shape (..., seq); x's head axes are unsqueezed in from the rank difference, so (seq,) and (batch, seq) both align with (batch, n_head, seq, head_dim).
Both pairing conventions are props, since they are a permutation of the feature axis apart and the choice is fixed by the weights: "half" pairs i with i + d/2 (HF rotate_half), "adjacent" pairs 2i with 2i + 1 (torchtune). Default "half".
Linear position interpolation (arXiv:2306.15595) and static NTK-aware scaling as props.
head_dim need not be static — the frequency ladder is built with pt.arange and folds to one constant when the size is known.
Everything is computed in x.type.dtype, so no float64 leaks into single-precision graphs.
Composes with AttentionLayer unchanged — _sdpa_graph is already position-agnostic.

base.py gains PositionalLayer, a sibling ABC for layers that take positions. Sibling rather than subclass because narrowing Layer.__call__ from one tensor to two is not a valid override; I checked, and subclassing still errors under mypy, so it would only move the type: ignore up a level. There is now no ignore at all.

Correction

An earlier version of this description claimed SplitDims/JoinDims have no JAX or MLX conversion, and that MultiheadAttention is therefore not JAX-compilable. Both claims were wrong — thanks @jessegrabowski for catching it. They are rewritten to Reshape during canonicalization and compile fine.

The real defect was in our own harness: tests/dispatch/jax/test_basic.py::compare_jax_and_py built its own RewriteDatabaseQuery(include=["jax"], ...), which omits canonicalize, so any op that only reaches the backend after being rewritten raises No JAX conversion there while compiling fine in real use. It now uses JAX._optimizer. Fixed in this PR — details on the _join_pairs thread.

Evidence

365 passed, 3 skipped (baseline on main: 312 passed, 3 skipped). mypy clean, pre-commit clean, CI green on Linux and Windows.

Correctness is pinned against an independent scalar reference in tests/test_positional.py that imports neither pytensor nor pytensor_ml, indexes scalars instead of slicing halves, and applies each 2×2 rotation literally — so a shared mistake in pairing, sign or frequency ladder cannot cancel out. Flipping the pairing on the reference side is asserted to fail, so the test cannot pass on a flipped implementation.

Also covered: step-by-step rotation equals the whole-sequence result at every step (the decode invariant); scores depend only on relative position; rotation is orthogonal and position 0 is the identity; linear scaling by 4 at position 8 equals plain RoPE at position 2; a symbolic feature axis matches the reference; RMSNorm leaves the mean intact where LayerNorm removes it; rewrite_for_prediction is a no-op for both, matching the LayerNorm precedent; gradients reach every channel through both pairings.

tests/dispatch/jax/test_positional.py confirms both ops lower on JAX through pytensor's OpFromGraph fallback with no jax_funcify of their own. No MLX test: mlx is not installable in the environment I verified in, and I am not shipping test code I could not run.

Two additive decoder primitives, from the in-scope set agreed on pymc-devs#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.

codecov-commenter · 2026-08-07T20:43:36Z

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered ☂️

jessegrabowski · 2026-08-08T03:19:35Z

    return (X - mu) / pt.sqrt(sigma_sq + epsilon), mu, sigma_sq


+def _rms_normalize(X, epsilon):


This is only called once, just inline it

jessegrabowski · 2026-08-08T03:20:29Z

    return inferred


+def _scale_parameter(name: str, n_in: int) -> TrainableParameter:


I don't like this, it makes _affine_parameters asymmetric. At this point just inline the variable creation everywhere.

jessegrabowski · 2026-08-08T03:21:13Z

+        if not self.affine:
+            return [X_normalized]
+
+        # Scale-only, so the affine transform contributes one input rather than the ``(loc, scale)``


jessegrabowski · 2026-08-08T03:22:10Z

+
+        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


Simplify this, it's weirdly obsessed with the fact that there's no location, so what. Nice to credit who we're copying and where it's used, though.

jessegrabowski · 2026-08-08T03:22:34Z

+        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,


simplify, don't need the huge essay

also i'd rather we just pick something for our library

jessegrabowski · 2026-08-08T03:35:48Z

+        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


needs References section and correct cross-reference for academic paper

jessegrabowski · 2026-08-08T03:36:18Z

+    Apply this to queries and keys before :func:`~pytensor_ml.layers.attention.scaled_dot_product_attention`,
+    never to values.
+
+    Parameters


This all needs to be dialed in, each parameter has the book of genesis under it

jessegrabowski · 2026-08-08T03:36:50Z

+    scaling_factor : float, optional
+        Extension factor for the scaled variants. Default is 1.0.
+
+    See Also


I don't personally care for the see also section, ymmv

jessegrabowski · 2026-08-08T03:38:21Z

+        self.scaling = scaling
+        self.scaling_factor = scaling_factor
+
+    # Positions are a required second input, which the one-tensor Layer.__call__ signature does not


So create a new subclass

jessegrabowski · 2026-08-08T03:40:16Z

-    # 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


these changes look like they were mixed in from other PRs

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.

cetagostini · 2026-08-08T15:21:05Z

hey jesse, addressed your comments.

activations.py/constant_like churn reverted, file is untouched again. epsilon is pt.constant(..., dtype=X.type.dtype) now
_rms_normalize and _scale_parameter inlined, _affine_parameters symmetric again, dead comments gone
static head_dim dropped, freqs built with pt.arange and fold to a constant, all in x.type.dtype so no float64
you were right about split/join. the actual bug was ours: compare_jax_and_py omits canonicalize, so it reported a gap that doesn't exist. fixed
subclassing Layer still fails mypy, can't narrow __call__, so PositionalLayer is a sibling ABC. no type: ignore left

spelling/refs/permalinks/doc trimming done. 365 passed, mypy and pre-commit clean.

jessegrabowski

very nice, one more round and we're done i think

jessegrabowski · 2026-08-08T18:03:55Z

    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: ...


+class PositionalLayer(ABC):


Do do envision a lot of layers that need this? I would prefer a simple VariadicLayerOp with __call__(self, *inputs: pt.TensorLike) -> pt.TensorVariable to allow as many inputs as needed for non-trivial layers. I'm trying to think about possible usages, like RNN layers will want to be BinaryLayerOp with input and hidden_layer, so that's 2 like yours. Maybe that's all we need?

jessegrabowski · 2026-08-08T18:07:01Z

+
+    # 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)


absolutely not an ask for this PR, but we might have a win here from a custom zip_concatenate primitive here to do just one allocation instead of two if the models you care about have a lot of these Ops.

jessegrabowski · 2026-08-08T18:14:56Z

+        )
+        angles = position_ids[..., None].astype(dtype) * inverse_frequencies
+        angles = _add_head_axes(angles, x, position_ids)
+        cos, sin = pt.cos(angles), pt.sin(angles)


nit: split this into two lines

jessegrabowski · 2026-08-08T18:15:26Z

+        cos, sin = pt.cos(angles), pt.sin(angles)
+
+        first, second = _split_pairs(x, self.pairing, head_dim)
+        rotated = _join_pairs(


nit: in general i prefer function calls to use keyword arguments for readability. In this case you also have expression arguments, it would read much more clearly with a line break between each argument.

jessegrabowski · 2026-08-08T18:37:06Z

+    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


These are just tests of pytensor though. We should just be running ci with a jax matrix value if we want to test in jax specifically (mlx too ideally). tests/dispatch/jax should be reserved for actual dispatches that we're adding.

jessegrabowski · 2026-08-08T18:37:41Z

+    np.testing.assert_allclose(res, expected, rtol=1e-5)
+
+
+def test_rms_norm_has_no_shift_parameter():


This can never error, it's not worth a test

I'd do a usefulness audit of all the tests to make sure each one can catch some kind of implementation or correctness drift in the future.

jessegrabowski · 2026-08-08T18:40:58Z

+
+
+def test_rope_pairing_conventions_are_not_interchangeable(rng):
+    # The pairing is the easiest thing to get silently wrong -- both conventions produce a plausible


most of the tests have comments under the definition... that's called a docstring. Also the test names are quite nice, do any of these comments tell me something the name doesn't?

jessegrabowski requested changes Aug 8, 2026

View reviewed changes

jessegrabowski added layer New or improvement to existing LayerOp enhancement New feature or request labels Aug 8, 2026

		return (X - mu) / pt.sqrt(sigma_sq + epsilon), mu, sigma_sq


		def _rms_normalize(X, epsilon):

		return inferred


		def _scale_parameter(name: str, n_in: int) -> TrainableParameter:


		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

		def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: ...


		class PositionalLayer(ABC):

		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

		np.testing.assert_allclose(res, expected, rtol=1e-5)


		def test_rms_norm_has_no_shift_parameter():



		def test_rope_pairing_conventions_are_not_interchangeable(rng):
		# The pairing is the easiest thing to get silently wrong -- both conventions produce a plausible

Uh oh!

Conversation

cetagostini commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

RMSNorm

RoPE

Correction

Evidence

Uh oh!

codecov-commenter commented Aug 7, 2026

Welcome to Codecov 🎉

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cetagostini commented Aug 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cetagostini commented Aug 7, 2026 •

edited

Loading

cetagostini commented Aug 8, 2026 •

edited

Loading