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/xtensor/basic.py

@@ -17,7 +17,7 @@ def perform(self, node, inputs, outputs):
     def do_constant_folding(self, fgraph, node):
         return False
 
-    def vectorize_node(self, node, *new_inputs):
+    def vectorize_node(self, node, *new_inputs, new_dim: str | None):
         raise NotImplementedError(f"Vectorized node not implemented for {self}")
 
 
@@ -30,7 +30,7 @@ class XTypeCastOp(TypeCastingOp):
     def infer_shape(self, fgraph, node, input_shapes):
         return input_shapes
 
-    def vectorize_node(self, node, *new_inputs):
+    def vectorize_node(self, node, *new_inputs, new_dim: str | None):
         raise NotImplementedError(f"Vectorized node not implemented for {self}")
 
 
@@ -48,11 +48,11 @@ 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
         # We transpose batch dims to the left, for consistency with tensor vectorization
         new_x = new_x.transpose(..., *old_x.dims)
-        return self.make_node(new_x)
+        return [self(new_x)]
 
 
 tensor_from_xtensor = TensorFromXTensor()
@@ -75,15 +75,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:
-            # TODO: Figure out API for this?
-            raise NotImplementedError(
-                f"Vectorization of {self} with batched inputs not implemented, "
-                "as it can't infer new dimension labels"
-            )
-        return self().make_node(new_x)
+            if new_dim is None:
+                raise NotImplementedError(
+                    f"Vectorization of {self} is not well defined because it can't infer the new dimension labels. "
+                    f"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):
@@ -107,15 +109,15 @@ 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))
 
         # new_dims may include a mix of old dims (possibly re-ordered), and new dims which won't be renamed
         new_dims = tuple(
             old_dim_mapping.get(new_dim, new_dim) for new_dim in new_x.dims
         )
-        return type(self)(new_dims).make_node(new_x)
+        return [type(self)(new_dims)(new_x)]
 
 
 def rename(x, name_dict: dict[str, str] | None = None, **names: str):

pytensor/xtensor/reduction.py

@@ -46,8 +46,8 @@ 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):
-        return self.make_node(new_x)
+    def vectorize_node(self, node, new_x, new_dim):
+        return [self(new_x)]
 
 
 def _process_user_dims(x: XTensorVariable, dim: REDUCE_DIM) -> Sequence[str]:
@@ -120,8 +120,8 @@ def make_node(self, x):
         out = x.type()
         return Apply(self, [x], [out])
 
-    def vectorize_node(self, node, new_x):
-        return self.make_node(new_x)
+    def vectorize_node(self, node, new_x, new_dim):
+        return [self(new_x)]
 
 
 def cumreduce(x, dim: REDUCE_DIM, *, binary_op):

pytensor/xtensor/shape.py

@@ -68,8 +68,8 @@ def make_node(self, x):
         )
         return Apply(self, [x], [output])
 
-    def vectorize_node(self, node, new_x):
-        return self.make_node(new_x)
+    def vectorize_node(self, node, new_x, new_dim):
+        return [self(new_x)]
 
 
 def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str]):
@@ -149,18 +149,13 @@ def make_node(self, x, *unstacked_length):
         )
         return Apply(self, [x, *unstacked_lengths], [output])
 
-    def vectorize_node(self, node, new_x, *new_unstacked_length):
-        if len(new_unstacked_length) != len(self.unstacked_dims):
-            raise NotImplementedError(
-                f"Vectorization of {self} with additional unstacked_length not implemented, "
-                "as it can't infer new dimension labels"
-            )
+    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(
                 f"Vectorization of {self} with batched unstacked_length not implemented, "
             )
-        return self.make_node(new_x, *new_unstacked_length)
+        return [self(new_x, *new_unstacked_length)]
 
 
 def unstack(x, dim: dict[str, dict[str, int]] | None = None, **dims: dict[str, int]):
@@ -205,10 +200,10 @@ 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)).make_node(new_x)
+        return [type(self)(dims=(*new_dims, *old_dims))(new_x)]
 
 
 def transpose(
@@ -323,8 +318,8 @@ 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):
-        return self.make_node(*new_inputs)
+    def vectorize_node(self, node, *new_inputs, new_dim):
+        return [self(*new_inputs)]
 
 
 def concat(xtensors, dim: str):
@@ -407,8 +402,8 @@ def make_node(self, x):
         )
         return Apply(self, [x], [out])
 
-    def vectorize_node(self, node, new_x):
-        return self.make_node(new_x)
+    def vectorize_node(self, node, new_x, new_dim):
+        return [self(new_x)]
 
 
 def squeeze(x, dim: str | Sequence[str] | None = None):
@@ -469,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(
@@ -572,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 := (
@@ -583,7 +578,7 @@ def vectorize_node(self, node, *new_inputs):
                         f"has an excluded dimension {sorted(invalid_excluded)} that it did not have before."
                     )
 
-        return self.make_node(*new_inputs)
+        return self(*new_inputs, return_list=True)
 
 
 def broadcast(