WIP Basic labeled tensor functionality

TODO: Split Stack from commit

Author: ricardoV94

pytensor/xtensor/__init__.py

@@ -0,0 +1,12 @@
+import warnings
+
+import pytensor.xtensor.rewriting
+from pytensor.xtensor.type import (
+    XTensorType,
+    as_xtensor,
+    xtensor,
+    xtensor_constant,
+)
+
+
+warnings.warn("xtensor module is experimental and full of bugs")

pytensor/xtensor/basic.py

@@ -0,0 +1,84 @@
+from pytensor.graph import Apply, Op
+from pytensor.tensor.type import TensorType
+from pytensor.xtensor.type import XTensorType, as_xtensor, xtensor
+
+
+class XOp(Op):
+    """A base class for XOps that shouldn't be materialized"""
+
+    def perform(self, node, inputs, outputs):
+        raise NotImplementedError(
+            f"xtensor operation {self} must be lowered to equivalent tensor operations"
+        )
+
+
+class XViewOp(Op):
+    # Make this a View Op with C-implementation
+    view_map = {0: [0]}
+
+    def perform(self, node, inputs, output_storage):
+        output_storage[0][0] = inputs[0]
+
+
+class TensorFromXTensor(XViewOp):
+    __props__ = ()
+
+    def make_node(self, x) -> Apply:
+        if not isinstance(x.type, XTensorType):
+            raise TypeError(f"x must be have an XTensorType, got {type(x.type)}")
+        output = TensorType(x.type.dtype, shape=x.type.shape)()
+        return Apply(self, [x], [output])
+
+
+tensor_from_xtensor = TensorFromXTensor()
+
+
+class XTensorFromTensor(XViewOp):
+    __props__ = ("dims",)
+
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+
+    def make_node(self, x) -> Apply:
+        if not isinstance(x.type, TensorType):
+            raise TypeError(f"x must be an TensorType type, got {type(x.type)}")
+        output = xtensor(dtype=x.type.dtype, dims=self.dims, shape=x.type.shape)
+        return Apply(self, [x], [output])
+
+
+def xtensor_from_tensor(x, dims):
+    return XTensorFromTensor(dims=dims)(x)
+
+
+class Rename(XViewOp):
+    __props__ = ("new_dims",)
+
+    def __init__(self, new_dims: tuple[str, ...]):
+        super().__init__()
+        self.new_dims = new_dims
+
+    def make_node(self, x):
+        x = as_xtensor(x)
+        output = x.type.clone(dims=self.new_dims)()
+        return Apply(self, [x], [output])
+
+
+def rename(x, name_dict: dict[str, str] | None = None, **names: str):
+    if name_dict is not None:
+        if names:
+            raise ValueError("Cannot use both positional and keyword names in rename")
+        names = name_dict
+
+    x = as_xtensor(x)
+    old_names = x.type.dims
+    new_names = list(old_names)
+    for old_name, new_name in names.items():
+        try:
+            new_names[old_names.index(old_name)] = new_name
+        except IndexError:
+            raise ValueError(
+                f"Cannot rename {old_name} to {new_name}: {old_name} not in {old_names}"
+            )
+
+    return Rename(tuple(new_names))(x)

pytensor/xtensor/rewriting/__init__.py

@@ -0,0 +1,2 @@
+import pytensor.xtensor.rewriting.basic
+import pytensor.xtensor.rewriting.shape

pytensor/xtensor/rewriting/basic.py

@@ -0,0 +1,49 @@
+from pytensor.graph import node_rewriter
+from pytensor.xtensor.basic import (
+    Rename,
+    TensorFromXTensor,
+    XTensorFromTensor,
+)
+from pytensor.xtensor.rewriting.utils import register_xcanonicalize
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[TensorFromXTensor])
+def useless_tensor_from_xtensor(fgraph, node):
+    """TensorFromXTensor(XTensorFromTensor(x)) -> x"""
+    [x] = node.inputs
+    if x.owner and isinstance(x.owner.op, XTensorFromTensor):
+        return [x.owner.inputs[0]]
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[XTensorFromTensor])
+def useless_xtensor_from_tensor(fgraph, node):
+    """XTensorFromTensor(TensorFromXTensor(x)) -> x"""
+    [x] = node.inputs
+    if x.owner and isinstance(x.owner.op, TensorFromXTensor):
+        return [x.owner.inputs[0]]
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[TensorFromXTensor])
+def useless_tensor_from_xtensor_of_rename(fgraph, node):
+    """TensorFromXTensor(Rename(x)) -> TensorFromXTensor(x)"""
+    [renamed_x] = node.inputs
+    if renamed_x.owner and isinstance(renamed_x.owner.op, Rename):
+        [x] = renamed_x.owner.inputs
+        return node.op(x, return_list=True)
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[Rename])
+def useless_rename(fgraph, node):
+    """
+
+    Rename(Rename(x, inner_dims), outer_dims) -> Rename(x, outer_dims)
+    Rename(X, XTensorFromTensor(x, inner_dims), outer_dims) -> XTensorFrom_tensor(x, outer_dims)
+    """
+    [renamed_x] = node.inputs
+    if renamed_x.owner and isinstance(renamed_x.owner.op, Rename | XTensorFromTensor):
+        [x] = renamed_x.owner.inputs
+        return node.op(x, return_list=True)

pytensor/xtensor/rewriting/shape.py

@@ -0,0 +1,29 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor import moveaxis
+from pytensor.xtensor.basic import tensor_from_xtensor, xtensor_from_tensor
+from pytensor.xtensor.rewriting.basic import register_xcanonicalize
+from pytensor.xtensor.shape import Stack
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[Stack])
+def lower_stack(fgraph, node):
+    [x] = node.inputs
+    batch_ndim = x.type.ndim - len(node.op.stacked_dims)
+    stacked_axes = [
+        i for i, dim in enumerate(x.type.dims) if dim in node.op.stacked_dims
+    ]
+    end = tuple(range(-len(stacked_axes), 0))
+
+    x_tensor = tensor_from_xtensor(x)
+    x_tensor_transposed = moveaxis(x_tensor, source=stacked_axes, destination=end)
+    if batch_ndim == (x.type.ndim - 1):
+        # This happens when we stack a "single" dimension, in this case all we need is the transpose
+        # Note: If we have meaningful rewrites before lowering, consider canonicalizing this as a Transpose + Rename
+        final_tensor = x_tensor_transposed
+    else:
+        final_shape = (*tuple(x_tensor_transposed.shape)[:batch_ndim], -1)
+        final_tensor = x_tensor_transposed.reshape(final_shape)
+
+    new_out = xtensor_from_tensor(final_tensor, dims=node.outputs[0].type.dims)
+    return [new_out]

pytensor/xtensor/rewriting/utils.py

@@ -0,0 +1,33 @@
+from pytensor.compile import optdb
+from pytensor.graph.rewriting.basic import NodeRewriter
+from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase
+
+
+optdb.register(
+    "xcanonicalize",
+    EquilibriumDB(ignore_newtrees=False),
+    "fast_run",
+    "fast_compile",
+    "xtensor",
+    position=0,
+)
+
+
+def register_xcanonicalize(
+    node_rewriter: RewriteDatabase | NodeRewriter | str, *tags: str, **kwargs
+):
+    if isinstance(node_rewriter, str):
+
+        def register(inner_rewriter: RewriteDatabase | NodeRewriter):
+            return register_xcanonicalize(
+                inner_rewriter, node_rewriter, *tags, **kwargs
+            )
+
+        return register
+
+    else:
+        name = kwargs.pop("name", None) or node_rewriter.__name__
+        optdb["xtensor"].register(
+            name, node_rewriter, "fast_run", "fast_compile", *tags, **kwargs
+        )
+        return node_rewriter

pytensor/xtensor/shape.py

@@ -0,0 +1,71 @@
+from collections.abc import Sequence
+
+from pytensor.graph import Apply
+from pytensor.xtensor.basic import XOp
+from pytensor.xtensor.type import as_xtensor, xtensor
+
+
+class Stack(XOp):
+    __props__ = ("new_dim_name", "stacked_dims")
+
+    def __init__(self, new_dim_name: str, stacked_dims: tuple[str, ...]):
+        super().__init__()
+        if new_dim_name in stacked_dims:
+            raise ValueError(
+                f"Stacking dim {new_dim_name} must not be in {stacked_dims}"
+            )
+        if not stacked_dims:
+            raise ValueError(f"Stacking dims must not be empty: got {stacked_dims}")
+        self.new_dim_name = new_dim_name
+        self.stacked_dims = stacked_dims
+
+    def make_node(self, x):
+        x = as_xtensor(x)
+        if not (set(self.stacked_dims) <= set(x.type.dims)):
+            raise ValueError(
+                f"Stacking dims {self.stacked_dims} must be a subset of {x.type.dims}"
+            )
+        if self.new_dim_name in x.type.dims:
+            raise ValueError(
+                f"Stacking dim {self.new_dim_name} must not be in {x.type.dims}"
+            )
+        if len(self.stacked_dims) == x.type.ndim:
+            batch_dims, batch_shape = (), ()
+        else:
+            batch_dims, batch_shape = zip(
+                *(
+                    (dim, shape)
+                    for dim, shape in zip(x.type.dims, x.type.shape)
+                    if dim not in self.stacked_dims
+                )
+            )
+        stack_shape = 1
+        for dim, shape in zip(x.type.dims, x.type.shape):
+            if dim in self.stacked_dims:
+                if shape is None:
+                    stack_shape = None
+                    break
+                else:
+                    stack_shape *= shape
+        output = xtensor(
+            dtype=x.type.dtype,
+            shape=(*batch_shape, stack_shape),
+            dims=(*batch_dims, self.new_dim_name),
+        )
+        return Apply(self, [x], [output])
+
+
+def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str]):
+    if dim is not None:
+        if dims:
+            raise ValueError("Cannot use both positional dim and keyword dims in stack")
+        dims = dim
+
+    y = x
+    for new_dim_name, stacked_dims in dims.items():
+        if isinstance(stacked_dims, str):
+            raise TypeError(
+                f"Stacking dims must be a sequence of strings, got a single string: {stacked_dims}"
+            )
+        y = Stack(new_dim_name, tuple(stacked_dims))(y)
+    return y