Implement vectorize_node for XOps

Author: ricardoV94

pyproject.toml

@@ -163,7 +163,7 @@ lines-after-imports = 2
 "tests/link/numba/**/test_*.py" = ["E402"]
 "tests/link/pytorch/**/test_*.py" = ["E402"]
 "tests/link/mlx/**/test_*.py" = ["E402"]
-"tests/xtensor/**/test_*.py" = ["E402"]
+"tests/xtensor/**/*.py" = ["E402"]
 
 
 

pytensor/xtensor/basic.py

@@ -17,6 +17,9 @@ def perform(self, node, inputs, outputs):
     def do_constant_folding(self, fgraph, node):
         return False
 
+    def vectorize_node(self, node, *new_inputs):
+        raise NotImplementedError(f"Vectorized node not implemented for {self}")
+
 
 class XTypeCastOp(TypeCastingOp):
     """Base class for Ops that type cast between TensorType and XTensorType.
@@ -27,6 +30,9 @@ class XTypeCastOp(TypeCastingOp):
     def infer_shape(self, fgraph, node, input_shapes):
         return input_shapes
 
+    def vectorize_node(self, node, *new_inputs):
+        raise NotImplementedError(f"Vectorized node not implemented for {self}")
+
 
 class TensorFromXTensor(XTypeCastOp):
     __props__ = ()
@@ -42,6 +48,12 @@ 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):
+        [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)
+
 
 tensor_from_xtensor = TensorFromXTensor()
 
@@ -63,6 +75,16 @@ 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):
+        [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)
+
 
 def xtensor_from_tensor(x, dims, name=None):
     return XTensorFromTensor(dims=dims)(x, name=name)
@@ -85,6 +107,16 @@ 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):
+        [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)
+
 
 def rename(x, name_dict: dict[str, str] | None = None, **names: str):
     if name_dict is not None:

pytensor/xtensor/indexing.py

@@ -4,13 +4,15 @@
 # https://numpy.org/neps/nep-0021-advanced-indexing.html
 # https://docs.xarray.dev/en/latest/user-guide/indexing.html
 # https://tutorial.xarray.dev/intermediate/indexing/advanced-indexing.html
+from itertools import chain
 from typing import Literal
 
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.scalar.basic import discrete_dtypes
 from pytensor.tensor.basic import as_tensor
 from pytensor.tensor.type_other import NoneTypeT, SliceType, make_slice
 from pytensor.xtensor.basic import XOp, xtensor_from_tensor
+from pytensor.xtensor.shape import broadcast
 from pytensor.xtensor.type import XTensorType, as_xtensor, xtensor
 
 
@@ -195,6 +197,15 @@ 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):
+        # 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)
+        old_x, *_ = node.inputs
+        dims_to_idxs = dict(zip(old_x.dims, new_idxs, strict=True))
+        new_idxs = tuple(dims_to_idxs.get(dim, slice(None)) for dim in new_x.dims)
+        return self.make_node(new_x, *new_idxs)
+
 
 index = Index()
 
@@ -226,6 +237,22 @@ def make_node(self, x, y, *idxs):
         out = x.type()
         return Apply(self, [x, y, *idxs], [out])
 
+    def vectorize_node(self, node, *new_inputs):
+        # If y or the indices have new dimensions we need to broadcast_x
+        exclude: set[str] = set(
+            chain.from_iterable(old_inp.dims for old_inp in node.inputs)
+        )
+        for new_inp, old_inp in zip(new_inputs, node.inputs, strict=True):
+            # Note: This check may be too conservative
+            if invalid_new_dims := ((set(new_inp.dims) - set(old_inp.dims)) & exclude):
+                raise NotImplementedError(
+                    f"Vectorize of {self} is undefined because one of the inputs {new_inp} new dimensions "
+                    f"was present in the old inputs: {sorted(invalid_new_dims)}"
+                )
+        new_x, *_ = broadcast(*new_inputs, exclude=tuple(exclude))
+        _, new_y, *new_idxs = new_inputs
+        return self.make_node(new_x, new_y, *new_idxs)
+
 
 index_assignment = IndexUpdate("set")
 index_increment = IndexUpdate("inc")

pytensor/xtensor/reduction.py

@@ -46,6 +46,9 @@ 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 _process_user_dims(x: XTensorVariable, dim: REDUCE_DIM) -> Sequence[str]:
     if isinstance(dim, str):
@@ -117,6 +120,9 @@ 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 cumreduce(x, dim: REDUCE_DIM, *, binary_op):
     x = as_xtensor(x)

pytensor/xtensor/shape.py

@@ -68,6 +68,9 @@ def make_node(self, x):
         )
         return Apply(self, [x], [output])
 
+    def vectorize_node(self, node, new_x):
+        return self.make_node(new_x)
+
 
 def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str]):
     if dim is not None:
@@ -146,6 +149,19 @@ 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"
+            )
+        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)
+
 
 def unstack(x, dim: dict[str, dict[str, int]] | None = None, **dims: dict[str, int]):
     if dim is not None:
@@ -189,6 +205,11 @@ def make_node(self, x):
         )
         return Apply(self, [x], [output])
 
+    def vectorize_node(self, node, new_x):
+        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)
+
 
 def transpose(
     x,
@@ -302,6 +323,9 @@ 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 concat(xtensors, dim: str):
     """Concatenate a sequence of XTensorVariables along a specified dimension.
@@ -383,6 +407,9 @@ def make_node(self, x):
         )
         return Apply(self, [x], [out])
 
+    def vectorize_node(self, node, new_x):
+        return self.make_node(new_x)
+
 
 def squeeze(x, dim: str | Sequence[str] | None = None):
     """Remove dimensions of size 1 from an XTensorVariable."""
@@ -442,6 +469,14 @@ def make_node(self, x, size):
         )
         return Apply(self, [x, size], [out])
 
+    def vectorize_node(self, node, new_x, new_size):
+        new_size = new_size.squeeze()
+        if new_size.type.ndim != 0:
+            raise NotImplementedError(
+                f"Vectorization of {self} with batched new_size not implemented, "
+            )
+        return self.make_node(new_x, new_size)
+
 
 def expand_dims(x, dim=None, axis=None, **dim_kwargs):
     """Add one or more new dimensions to an XTensorVariable."""
@@ -537,6 +572,19 @@ def make_node(self, *inputs):
 
         return Apply(self, inputs, outputs)
 
+    def vectorize_node(self, node, *new_inputs):
+        if exclude_set := set(self.exclude):
+            for new_x, old_x in zip(node.inputs, new_inputs, strict=True):
+                if invalid_excluded := (
+                    (set(new_x.dims) - set(old_x.dims)) & exclude_set
+                ):
+                    raise NotImplementedError(
+                        f"Vectorize of {self} is undefined because one of the inputs {new_x} "
+                        f"has an excluded dimension {sorted(invalid_excluded)} that it did not have before."
+                    )
+
+        return self.make_node(*new_inputs)
+
 
 def broadcast(
     *args, exclude: str | Sequence[str] | None = None

pytensor/xtensor/type.py

@@ -1044,7 +1044,7 @@ def as_xtensor(x, dims: Sequence[str] | None = None, *, name: str | None = None)
 
     if isinstance(x, Variable):
         if isinstance(x.type, XTensorType):
-            if (dims is None) or (x.type.dims == dims):
+            if (dims is None) or (x.type.dims == tuple(dims)):
                 return x
             else:
                 raise ValueError(

pytensor/xtensor/vectorization.py

@@ -6,6 +6,7 @@
 from pytensor import scalar as ps
 from pytensor import shared
 from pytensor.graph import Apply, Op
+from pytensor.graph.replace import _vectorize_node
 from pytensor.scalar import discrete_dtypes
 from pytensor.tensor import tensor
 from pytensor.tensor.random.op import RNGConsumerOp
@@ -14,8 +15,11 @@
     get_static_shape_from_size_variables,
 )
 from pytensor.utils import unzip
-from pytensor.xtensor.basic import XOp
-from pytensor.xtensor.type import XTensorVariable, as_xtensor, xtensor
+from pytensor.xtensor.basic import (
+    XOp,
+    XTypeCastOp,
+)
+from pytensor.xtensor.type import XTensorType, XTensorVariable, as_xtensor, xtensor
 
 
 def combine_dims_and_shape(
@@ -69,6 +73,9 @@ def make_node(self, *inputs):
         ]
         return Apply(self, inputs, outputs)
 
+    def vectorize_node(self, node, *new_inputs):
+        return self.make_node(*new_inputs)
+
 
 class XBlockwise(XOp):
     __props__ = ("core_op", "core_dims")
@@ -136,6 +143,9 @@ def make_node(self, *inputs):
         ]
         return Apply(self, inputs, outputs)
 
+    def vectorize_node(self, node, *new_inputs):
+        return self.make_node(*new_inputs)
+
 
 class XRV(XOp, RNGConsumerOp):
     """Wrapper for RandomVariable operations that follows xarray-like broadcasting semantics.
@@ -283,3 +293,60 @@ def make_node(self, rng, *extra_dim_lengths_and_params):
         )
 
         return Apply(self, [rng, *extra_dim_lengths, *params], [rng.type(), out])
+
+    def vectorize_node(self, node, *new_inputs):
+        if len(new_inputs) != len(node.inputs):
+            # TODO: Figure out API to allow this
+            raise NotImplementedError(
+                f"Vectorization of {self} with additional extra_dim_lengths not implemented, "
+                "as it can't infer new dimension labels"
+            )
+        new_rng, *new_extra_dim_lengths_and_params = new_inputs
+        new_extra_dim_lengths, new_params = (
+            new_extra_dim_lengths_and_params[: len(self.extra_dims)],
+            new_extra_dim_lengths_and_params[len(self.extra_dims) :],
+        )
+
+        new_extra_dim_lengths = [dl.squeeze() for dl in new_extra_dim_lengths]
+        if not all(dl.type.ndim == 0 for dl in new_extra_dim_lengths):
+            raise NotImplementedError(
+                f"Vectorization of {self} with batched extra_dim_lengths not implemented, "
+            )
+
+        return self.make_node(new_rng, *new_extra_dim_lengths, *new_params)
+
+
+@_vectorize_node.register(XOp)
+@_vectorize_node.register(XTypeCastOp)
+def vectorize_xop(op: XOp, node, *new_inputs) -> Apply:
+    old_inp_dims = [
+        inp.dims for inp in node.inputs if isinstance(inp.type, XTensorType)
+    ]
+    old_out_dims = [
+        out.dims for out in node.outputs if isinstance(out.type, XTensorType)
+    ]
+    all_old_dims_set = set(chain.from_iterable((*old_inp_dims, old_out_dims)))
+
+    for new_inp, old_inp in zip(new_inputs, node.inputs, strict=True):
+        if not (
+            isinstance(new_inp.type, XTensorType)
+            and isinstance(old_inp.type, XTensorType)
+        ):
+            continue
+
+        old_dims_set = set(old_inp.dims)
+        new_dims_set = set(new_inp.dims)
+
+        # Validate that new inputs didn't drop pre-existing dims
+        if missing_dims := old_dims_set - new_dims_set:
+            raise ValueError(
+                f"Vectorized input {new_inp} is missing pre-existing dims: {sorted(missing_dims)}"
+            )
+        # Or have new dimensions that were already in the graph
+        if new_core_dims := ((new_dims_set - old_dims_set) & all_old_dims_set):
+            raise ValueError(
+                f"Vectorized input {new_inp} has new dimensions that were present in the original graph: {new_core_dims}"
+            )
+
+    # TODO: Once we stop having to return an Apply, automatically align batch_dimensions in the order they were first seen
+    return op.vectorize_node(node, *new_inputs)