Implement dim-aware vectorize_graph

Author: ricardoV94

pytensor/graph/replace.py

@@ -283,6 +283,13 @@ def vectorize_graph(
         # [array([-10., -11.]), array([10., 11.])]
 
     """
+    # TODO: Move this to tensor.vectorize, and make this helper type agnostic.
+    #
+    # This helper may dispatch to tensor.vectorize_graph or xtensor.vectorize_graph depending on the replacement types
+    # The behavior is distinct, because tensor vectorization depends on axis-position while xtensor depends on dimension labels
+    #
+    # xtensor.vectorize_graph will be able to handle batched inner tensor operations, while tensor.vectorize_graph won't,
+    # as it is by design unaware of xtensors and their semantics.
     if isinstance(outputs, Sequence):
         seq_outputs = outputs
     else:

pytensor/tensor/extra_ops.py

@@ -1,5 +1,5 @@
 import warnings
-from collections.abc import Collection, Iterable
+from collections.abc import Collection, Iterable, Sequence
 from textwrap import dedent
 
 import numpy as np
@@ -1926,7 +1926,7 @@ def logspace(
 
 
 def broadcast_to(
-    x: TensorVariable, shape: TensorVariable | tuple[Variable, ...]
+    x: TensorLike, shape: TensorLike | Sequence[TensorLike]
 ) -> TensorVariable:
     """Broadcast an array to a new shape.
 

pytensor/xtensor/basic.py

@@ -18,7 +18,9 @@ def perform(self, node, inputs, outputs):
     def do_constant_folding(self, fgraph, node):
         return False
 
-    def vectorize_node(self, node, *new_inputs) -> Sequence[Variable]:
+    def vectorize_node(
+        self, node, *new_inputs, new_dim: str | None
+    ) -> Sequence[Variable]:
         raise NotImplementedError(f"Vectorized node not implemented for {self}")
 
 
@@ -31,7 +33,9 @@ class XTypeCastOp(TypeCastingOp):
     def infer_shape(self, fgraph, node, input_shapes):
         return input_shapes
 
-    def vectorize_node(self, node, *new_inputs) -> Sequence[Variable]:
+    def vectorize_node(
+        self, node, *new_inputs, new_dim: str | None
+    ) -> Sequence[Variable]:
         raise NotImplementedError(f"Vectorized node not implemented for {self}")
 
 
@@ -49,12 +53,13 @@ def L_op(self, inputs, outs, g_outs):
         [g_out] = g_outs
         return [xtensor_from_tensor(g_out, dims=x.type.dims)]
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         [old_x] = node.inputs
         if (new_x.ndim - old_x.ndim) > 1:
             raise NotImplementedError(
                 f"Vectorization of {self} cannot guarantee correct placement of multiple batch dimensions. "
-                "You can call vectorize_graph one batch dimension at a time."
+                "You can call vectorize_graph one batch dimension at a time, "
+                "or pytensor.xtensor.vectorization.vectorize_graph instead."
             )
         new_x = new_x.transpose(..., *old_x.dims)
         return [self(new_x)]
@@ -80,14 +85,17 @@ def L_op(self, inputs, outs, g_outs):
         [g_out] = g_outs
         return [tensor_from_xtensor(g_out)]
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         [old_x] = node.inputs
         if new_x.ndim != old_x.ndim:
-            raise NotImplementedError(
-                f"Vectorization of {self} with batched inputs not implemented, "
-                "as it can't infer new dimension labels"
-            )
-        return [self(new_x)]
+            if new_dim is None:
+                raise NotImplementedError(
+                    f"Vectorization of {self} cannot infer the new dimension labels. "
+                    "Use pytensor.xtensor.vectorization.vectorize_graph instead."
+                )
+            return [type(self)(dims=(new_dim, *self.dims))(new_x)]
+        else:
+            return [self(new_x)]
 
 
 def xtensor_from_tensor(x, dims, name=None):
@@ -111,7 +119,7 @@ def L_op(self, inputs, outs, g_outs):
         [g_out] = g_outs
         return [rename(g_out, dims=x.type.dims)]
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         [old_x] = node.inputs
         old_dim_mapping = dict(zip(old_x.dims, self.new_dims, strict=True))
 

pytensor/xtensor/indexing.py

@@ -197,7 +197,7 @@ def combine_dim_info(idx_dim, idx_dim_shape):
         output = xtensor(dtype=x.type.dtype, shape=out_shape, dims=out_dims)
         return Apply(self, [x, *idxs], [output])
 
-    def vectorize_node(self, node, new_x, *new_idxs):
+    def vectorize_node(self, node, new_x, *new_idxs, new_dim):
         # new_x may have dims in different order
         # we pair each pre-existing dim to the respective index
         # with new dims having simply a slice(None)
@@ -237,7 +237,7 @@ def make_node(self, x, y, *idxs):
         out = x.type()
         return Apply(self, [x, y, *idxs], [out])
 
-    def vectorize_node(self, node, *new_inputs):
+    def vectorize_node(self, node, *new_inputs, new_dim):
         # If y or the indices have new dimensions we need to broadcast_x
         exclude: set[str] = set(
             chain.from_iterable(

pytensor/xtensor/reduction.py

@@ -46,7 +46,7 @@ def make_node(self, x):
         output = xtensor(dtype=x.type.dtype, shape=out_shape, dims=out_dims)
         return Apply(self, [x], [output])
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         return [self(new_x)]
 
 
@@ -120,7 +120,7 @@ def make_node(self, x):
         out = x.type()
         return Apply(self, [x], [out])
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         return [self(new_x)]
 
 

pytensor/xtensor/shape.py

@@ -68,7 +68,7 @@ def make_node(self, x):
         )
         return Apply(self, [x], [output])
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         return [self(new_x)]
 
 
@@ -149,7 +149,7 @@ def make_node(self, x, *unstacked_length):
         )
         return Apply(self, [x, *unstacked_lengths], [output])
 
-    def vectorize_node(self, node, new_x, *new_unstacked_length):
+    def vectorize_node(self, node, new_x, *new_unstacked_length, new_dim):
         new_unstacked_length = [ul.squeeze() for ul in new_unstacked_length]
         if not all(ul.type.ndim == 0 for ul in new_unstacked_length):
             raise NotImplementedError(
@@ -200,7 +200,7 @@ def make_node(self, x):
         )
         return Apply(self, [x], [output])
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         old_dims = self.dims
         new_dims = tuple(dim for dim in new_x.dims if dim not in old_dims)
         return [type(self)(dims=(*new_dims, *old_dims))(new_x)]
@@ -318,7 +318,7 @@ def make_node(self, *inputs):
         output = xtensor(dtype=dtype, dims=dims, shape=shape)
         return Apply(self, inputs, [output])
 
-    def vectorize_node(self, node, *new_inputs):
+    def vectorize_node(self, node, *new_inputs, new_dim):
         return [self(*new_inputs)]
 
 
@@ -402,7 +402,7 @@ def make_node(self, x):
         )
         return Apply(self, [x], [out])
 
-    def vectorize_node(self, node, new_x):
+    def vectorize_node(self, node, new_x, new_dim):
         return [self(new_x)]
 
 
@@ -464,7 +464,7 @@ def make_node(self, x, size):
         )
         return Apply(self, [x, size], [out])
 
-    def vectorize_node(self, node, new_x, new_size):
+    def vectorize_node(self, node, new_x, new_size, new_dim):
         new_size = new_size.squeeze()
         if new_size.type.ndim != 0:
             raise NotImplementedError(
@@ -567,7 +567,7 @@ def make_node(self, *inputs):
 
         return Apply(self, inputs, outputs)
 
-    def vectorize_node(self, node, *new_inputs):
+    def vectorize_node(self, node, *new_inputs, new_dim):
         if exclude_set := set(self.exclude):
             for new_x, old_x in zip(node.inputs, new_inputs, strict=True):
                 if invalid_excluded := (