Update to pytensor>=3.2.3 (#40)

* Decouple graph serialization from FrozenFunctionGraph constructor

* Base layer ops on SymbolicOp so numba optimizes their inner graphs

* Bump pytensor lower pin to 3.2.3

Author: jessegrabowski

conda_envs/pytensor_ml-gpu_jax.yml

@@ -6,7 +6,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.0.0,<4.0.0
+  - pytensor>=3.2.3,<4.0.0
   - numpy
   - scikit-learn
 

conda_envs/pytensor_ml.yml

@@ -5,7 +5,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.0.0,<4.0.0
+  - pytensor>=3.2.3,<4.0.0
   - numpy
   - safetensors
   - scikit-learn

pyproject.toml

@@ -34,7 +34,7 @@ keywords = [
 ]
 
 dependencies = [
-  "pytensor>=3.0.0,<4.0.0",
+  "pytensor>=3.2.3,<4.0.0",
   "numpy",
   "safetensors"
 ]

pytensor_ml/attention.py

@@ -11,6 +11,9 @@
 class AttentionLayer(UnaryLayerOp):
     __props__ = ("is_causal", "scale")
 
+    def build_inner_graph(self, q, k, v, mask=None):
+        return [_sdpa_graph(q, k, v, mask, self.is_causal, self.scale)]
+
 
 def _repeat_kv(x: TensorVariable, n_head: int | None, n_kv_head: int | None) -> TensorVariable:
     """Broadcast ``n_kv_head`` key/value heads up to ``n_head`` query heads (grouped-query attention).
@@ -117,11 +120,7 @@ def scaled_dot_product_attention(
         mask = pt.as_tensor(mask)
         inputs.append(mask)
 
-    out = _sdpa_graph(q, k, v, mask, is_causal, scale)
-
     op = AttentionLayer(
-        inputs=inputs,
-        outputs=[out],
         name="ScaledDotProductAttention",
         is_causal=is_causal,
         scale=scale,
@@ -184,8 +183,14 @@ def __init__(
         self.out_proj = Linear(f"{self.name}_out_proj", n_head * self.head_dim, n_embd, bias)
 
     def _split_heads(self, x: pt.TensorVariable, n_head: int) -> pt.TensorVariable:
-        # (..., seq, n_head * head_dim) -> (..., n_head, seq, head_dim)
-        x = x.reshape((*tuple(x.shape[:-1]), n_head, self.head_dim))
+        # (..., seq, n_head * head_dim) -> (..., n_head, seq, head_dim). Keep statically known leading
+        # dims as literals so the reshaped tensor retains its full static shape -- a symbolic
+        # ``x.shape[i]`` would erase it. A static core shape is what lets the numba backend vectorize
+        # attention (e.g. under ``vectorize_graph``) instead of falling back to object mode.
+        lead = tuple(
+            size if size is not None else x.shape[i] for i, size in enumerate(x.type.shape[:-1])
+        )
+        x = x.reshape((*lead, n_head, self.head_dim))
         return x.swapaxes(-3, -2)
 
     def __call__(self, x: pt.TensorLike, mask: pt.TensorLike | None = None) -> pt.TensorVariable:

pytensor_ml/json_serialize.py

@@ -7,7 +7,6 @@
 import pytensor.tensor as pt
 
 from pytensor.graph.basic import Constant, Variable
-from pytensor.graph.fg import FrozenFunctionGraph
 from pytensor.graph.op import Op
 from pytensor.graph.traversal import io_toposort
 from pytensor.graph.type import Type
@@ -136,12 +135,10 @@ def _resolve_class(path: str):
     return getattr(importlib.import_module(module), name)
 
 
-def graph_to_json(frozen: FrozenFunctionGraph) -> dict:
-    """Serialize a frozen graph to a dict of input types, op nodes, and output references."""
-    nodes = io_toposort(frozen.inputs, frozen.outputs)
-    var_ref: dict[int, dict] = {
-        id(inp): {"input": index} for index, inp in enumerate(frozen.inputs)
-    }
+def graph_to_json(inputs: Sequence[Variable], outputs: Sequence[Variable]) -> dict:
+    """Serialize a graph to a dict of input types, op nodes, and output references."""
+    nodes = io_toposort(inputs, outputs)
+    var_ref: dict[int, dict] = {id(inp): {"input": index} for index, inp in enumerate(inputs)}
     for node_index, node in enumerate(nodes):
         for output_index, out in enumerate(node.outputs):
             var_ref[id(out)] = {"node": node_index, "out": output_index}
@@ -155,7 +152,7 @@ def make_ref(variable: Variable) -> dict:
         raise ValueError(f"Unresolved variable while serializing graph: {variable!r}")
 
     return {
-        "inputs": [type_to_json(inp.type) for inp in frozen.inputs],
+        "inputs": [type_to_json(inp.type) for inp in inputs],
         "nodes": [
             {
                 "op": op_to_json(node.op),
@@ -164,7 +161,7 @@ def make_ref(variable: Variable) -> dict:
             }
             for node in nodes
         ],
-        "outputs": [make_ref(out) for out in frozen.outputs],
+        "outputs": [make_ref(out) for out in outputs],
     }
 
 
@@ -221,8 +218,7 @@ def serialize_graph(inputs: Sequence[Variable], outputs: Sequence[Variable]) ->
     dict
         A JSON-native description of the graph's structure (no parameter values).
     """
-    frozen = FrozenFunctionGraph(list(inputs), list(outputs))
-    return graph_to_json(frozen)
+    return graph_to_json(list(inputs), list(outputs))
 
 
 def deserialize_graph(

pytensor_ml/layers.py

@@ -31,6 +31,12 @@ def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: ...
 class LinearLayer(UnaryLayerOp):
     __props__ = ("n_in", "n_out", "bias")
 
+    def build_inner_graph(self, X, W, b=None):
+        res = X @ W
+        if self.bias:
+            res = res + b
+        return [res]
+
 
 class Linear(Layer):
     def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
@@ -49,20 +55,15 @@ def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
     def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         X = pt.as_tensor(X)
 
-        init_st_shape = shape_to_str(X.type.shape)
-
-        res = X @ self.W
         inputs = [X, self.W]
         if self.bias:
-            res += self.b
             inputs.append(self.b)
 
-        final_st_shape = shape_to_str(res.type.shape)
+        input_shape = shape_to_str(X.type.shape)
+        output_shape = shape_to_str((X @ self.W).type.shape)
 
         ofg = LinearLayer(
-            inputs=inputs,
-            outputs=[res],
-            name=f"{self.name}[{init_st_shape} -> {final_st_shape}]",
+            name=f"{self.name}[{input_shape} -> {output_shape}]",
             n_in=self.n_in,
             n_out=self.n_out,
             bias=self.bias,
@@ -76,6 +77,9 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
 class EmbeddingLayer(UnaryLayerOp):
     __props__ = ("n_embeddings", "n_features")
 
+    def build_inner_graph(self, ids, W):
+        return [W[ids]]
+
 
 class Embedding(Layer):
     r"""
@@ -107,8 +111,6 @@ def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
         ids = pt.as_tensor(ids)
 
         out = EmbeddingLayer(
-            inputs=[ids, self.W],
-            outputs=[self.W[ids]],
             name=self.name,
             n_embeddings=self.n_embeddings,
             n_features=self.n_features,
@@ -121,6 +123,9 @@ def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
 class DropoutLayer(UnaryLayerOp):
     __props__ = ("p",)
 
+    def build_inner_graph(self, X, mask):
+        return [pt.where(mask, ift=X / (1 - self.p), iff=0)]
+
 
 class Dropout(Layer):
     def __init__(self, name: str | None = None, p: float = 0.5, random_state: Any | None = None):
@@ -137,8 +142,6 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         mask = mask.astype(config.floatX)
 
         X_masked = DropoutLayer(
-            inputs=[X, mask],
-            outputs=[pt.where(mask, ift=X / (1 - p), iff=0)],
             name=f"{self.name}[p = {self.p}]",
             p=self.p,
         )(X, mask)
@@ -147,20 +150,52 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         return X_masked
 
 
+def _batch_normalize(X, epsilon):
+    mu = X.mean(axis=0)
+    sigma_sq = X.var(axis=0)
+    return (X - mu) / pt.sqrt(sigma_sq + epsilon), mu, sigma_sq
+
+
 class BatchNormLayer(LayerOp):
     __props__ = ("n_in", "epsilon", "momentum", "affine")
 
     def update_map(self):
+        # Outputs 1 and 2 (the new running mean and variance) update inputs 3 and 4 (the old ones).
         return {1: 3, 2: 4}
 
+    def build_inner_graph(self, X, *rest):
+        if self.affine:
+            loc, scale, running_mean, running_var = rest
+        else:
+            running_mean, running_var = rest
+
+        X_normalized, mu, sigma_sq = _batch_normalize(X, self.epsilon)
+        X_rescaled = X_normalized * scale + loc if self.affine else X_normalized
+
+        new_running_mean = self.momentum * mu + (1 - self.momentum) * running_mean
+        new_running_var = self.momentum * sigma_sq + (1 - self.momentum) * running_var
+
+        return [X_rescaled, new_running_mean, new_running_var]
+
 
 class NoRunningStatsBatchNormLayer(UnaryLayerOp):
     __props__ = ("n_in", "epsilon", "momentum", "affine")
 
+    def build_inner_graph(self, X, *rest):
+        X_normalized, _, _ = _batch_normalize(X, self.epsilon)
+        if self.affine:
+            loc, scale = rest
+            return [X_normalized * scale + loc]
+        return [X_normalized]
+
 
 class PredictionBatchNormLayer(UnaryLayerOp):
     __props__ = ("n_in", "epsilon", "momentum", "affine")
 
+    def build_inner_graph(self, X, loc, scale, running_mean, running_var):
+        res = (X - running_mean) / pt.sqrt(running_var + self.epsilon)
+        return [loc + res * scale]
+
 
 class BatchNorm2D(Layer):
     r"""
@@ -261,26 +296,14 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
 
         self._initialize_params(X)
 
-        mu = X.mean(axis=0)
-        sigma_sq = X.var(axis=0)
-
-        X_normalized = (X - mu) / pt.sqrt(sigma_sq + self.epsilon)
-
         if self.affine:
             assert self.scale is not None and self.loc is not None
-            X_rescaled = X_normalized * self.scale + self.loc
             inputs.extend([self.loc, self.scale])
-        else:
-            X_rescaled = X_normalized
 
         if self.track_running_stats:
             assert self.running_mean is not None and self.running_var is not None
-            new_running_mean = self.momentum * mu + (1 - self.momentum) * self.running_mean
-            new_running_var = self.momentum * sigma_sq + (1 - self.momentum) * self.running_var
 
             batch_norm_op: LayerOp = BatchNormLayer(
-                inputs=[*inputs, self.running_mean, self.running_var],
-                outputs=[X_rescaled, new_running_mean, new_running_var],
                 name=self.name,
                 n_in=self.n_in,
                 epsilon=self.epsilon,
@@ -294,8 +317,6 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
 
         else:
             batch_norm_op = NoRunningStatsBatchNormLayer(
-                inputs=inputs,
-                outputs=[X_rescaled],
                 name=self.name,
                 n_in=self.n_in,
                 epsilon=self.epsilon,
@@ -315,6 +336,18 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
 class LayerNormLayer(UnaryLayerOp):
     __props__ = ("n_in", "epsilon", "affine")
 
+    def build_inner_graph(self, X, *rest):
+        mu = X.mean(axis=-1, keepdims=True)
+        # Biased variance (ddof=0), matching torch.nn.LayerNorm; pretrained weights assume it. Do
+        # not "correct" to the unbiased estimator -- it would diverge from every pretrained model.
+        sigma_sq = X.var(axis=-1, keepdims=True)
+        X_normalized = (X - mu) / pt.sqrt(sigma_sq + self.epsilon)
+
+        if self.affine:
+            scale, loc = rest
+            return [X_normalized * scale + loc]
+        return [X_normalized]
+
 
 class LayerNorm(Layer):
     r"""
@@ -380,23 +413,12 @@ def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         X = pt.as_tensor(X)
         self._initialize_params(X)
 
-        mu = X.mean(axis=-1, keepdims=True)
-        # Biased variance (ddof=0), matching torch.nn.LayerNorm; pretrained weights assume it. Do
-        # not "correct" to the unbiased estimator -- it would diverge from every pretrained model.
-        sigma_sq = X.var(axis=-1, keepdims=True)
-        X_normalized = (X - mu) / pt.sqrt(sigma_sq + self.epsilon)
-
         inputs = [X]
         if self.affine:
             assert self.scale is not None and self.loc is not None
-            X_rescaled = X_normalized * self.scale + self.loc
             inputs.extend([self.scale, self.loc])
-        else:
-            X_rescaled = X_normalized
 
         X_transformed = LayerNormLayer(
-            inputs=inputs,
-            outputs=[X_rescaled],
             name=self.name,
             n_in=self.n_in,
             epsilon=self.epsilon,

pytensor_ml/pytensorf.py

@@ -8,7 +8,7 @@
 
 from pytensor import Mode
 from pytensor.compile import Function, SharedVariable, get_mode
-from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.builders import OpFromGraph, SymbolicOp
 from pytensor.graph import FunctionGraph, RewriteDatabaseQuery, graph_inputs, rewrite_graph
 from pytensor.graph.basic import Apply, Constant, equal_computations
 from pytensor.graph.fg import Output
@@ -22,19 +22,18 @@
 RandomSeed = None | int | Sequence[int] | np.ndarray
 
 
-class LayerOp(OpFromGraph):
-    # This can be removed once https://github.com/pymc-devs/pytensor/issues/1114 is fixed
+class LayerOp(SymbolicOp):
+    """Base class for the library's neural-network ops.
 
-    __props__: tuple[str, ...] = ()
-
-    def __init__(self, *args, **kwargs):
-        prop_kwargs = {key: value for key, value in kwargs.items() if key in self.__props__}
-        kwargs = {key: value for key, value in kwargs.items() if key not in self.__props__}
-
-        for key, value in prop_kwargs.items():
-            setattr(self, key, value)
+    A ``SymbolicOp`` is an ``OpFromGraph`` whose inner graph is rebuilt from its ``__props__`` and input
+    types by :meth:`build_inner_graph`, so equal props with equal inputs yield an identical op. Basing the
+    layers on it (rather than plain ``OpFromGraph``) is what lets the numba backend optimize each inner
+    graph: ``SymbolicOp`` restores the fgraph-aware ``__eq__``/``__hash__`` that a props-carrying
+    ``OpFromGraph`` would otherwise lose, so the ``ofg_inner_graph`` rewrite keeps its optimized inner
+    graph instead of discarding it as unchanged.
+    """
 
-        super().__init__(*args, **kwargs)
+    __props__: tuple[str, ...] = ()
 
     def update_map(self) -> dict[int, int]:
         """Return a mapping of output indexes to input indexes"""
@@ -44,7 +43,7 @@ def update_map(self) -> dict[int, int]:
 class UnaryLayerOp(LayerOp):
     """A ``LayerOp`` with exactly one output, typed as such.
 
-    ``OpFromGraph.__call__`` is annotated ``Variable | list[Variable]`` because an op may produce many
+    ``SymbolicOp.__call__`` is annotated ``Variable | list[Variable]`` because an op may produce many
     outputs; a unary layer op produces one, so narrow the result to ``TensorVariable``. The ``isinstance``
     guard narrows for the type checker without a cast and asserts the invariant at runtime.
     """

pytensor_ml/rewriting/layers.py

@@ -1,5 +1,3 @@
-import pytensor.tensor as pt
-
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
@@ -60,12 +58,7 @@ def rewrite_batch_stats_to_running_average_stats(
     """
     X, loc, scale, running_mean, running_var = node.inputs
 
-    res = (X - running_mean) / pt.sqrt(running_var + node.op.epsilon)  # type: ignore[operator]
-    res = loc + res * scale
-
     batch_norm_op = PredictionBatchNormLayer(
-        inputs=[X, loc, scale, running_mean, running_var],
-        outputs=[res],
         name=f"{node.op.name}",
         n_in=node.op.n_in,
         momentum=node.op.momentum,