POC named tensors

Author: ricardoV94

pytensor/xtensor/__init__.py

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

pytensor/xtensor/basic.py

@@ -0,0 +1,73 @@
+from itertools import chain
+
+import pytensor.scalar as ps
+import pytensor.xtensor as px
+from pytensor.graph import Apply, Op
+from pytensor.tensor import TensorType
+
+
+class TensorFromXTensor(Op):
+    # TODO: May need mapping of named dims to positional dims?
+
+    def make_node(self, x) -> Apply:
+        if not isinstance(x.type, px.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])
+
+    def perform(self, node, inputs, output_storage) -> None:
+        [x] = inputs
+        output_storage[0][0] = x.copy()
+
+
+tensor_from_xtensor = TensorFromXTensor()
+
+
+class XTensorFromTensor(Op):
+    __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 = px.XTensorType(x.type.dtype, dims=self.dims, shape=x.type.shape)()
+        return Apply(self, [x], [output])
+
+    def perform(self, node, inputs, output_storage) -> None:
+        [x] = inputs
+        output_storage[0][0] = x.copy()
+
+
+def xtensor_from_tensor(x, dims):
+    return XTensorFromTensor(dims=dims)(x)
+
+
+class XElemwise(Op):
+    __props__ = ("scalar_op",)
+
+    def __init__(self, scalar_op):
+        super().__init__()
+        self.scalar_op = scalar_op
+
+    def make_node(self, *inputs):
+        # TODO: Check dim lengths match
+        inputs = [px.as_xtensor(inp) for inp in inputs]
+        # TODO: This ordering is different than what xarray does
+        unique_dims = sorted(set(chain.from_iterable(inp.type.dims for inp in inputs)))
+        # TODO: Fix dtype
+        output_type = px.XTensorType(
+            "float64", dims=unique_dims, shape=(None,) * len(unique_dims)
+        )
+        outputs = [output_type() for _ in range(self.scalar_op.nout)]
+        return Apply(self, inputs, outputs)
+
+    def perform(self, *args, **kwargs) -> None:
+        raise NotImplementedError(
+            "xtensor operations must be rewritten as tensor operations"
+        )
+
+
+add = XElemwise(ps.add)

pytensor/xtensor/rewriting/__init__.py

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

pytensor/xtensor/rewriting/basic.py

@@ -0,0 +1,30 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor import expand_dims
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.xtensor.basic import XElemwise, tensor_from_xtensor, xtensor_from_tensor
+from pytensor.xtensor.rewriting.utils import register_xcanonicalize
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[XElemwise])
+def xelemwise_to_elemwise(fgraph, node):
+    # Convert inputs to TensorVariables and add broadcastable dims
+    output_dims = node.outputs[0].type.dims
+
+    tensor_inputs = []
+    for inp in node.inputs:
+        inp_dims = inp.type.dims
+        axis = [i for i, dim in enumerate(output_dims) if dim not in inp_dims]
+        tensor_inp = tensor_from_xtensor(inp)
+        tensor_inp = expand_dims(tensor_inp, axis)
+        tensor_inputs.append(tensor_inp)
+
+    tensor_outs = Elemwise(scalar_op=node.op.scalar_op)(
+        *tensor_inputs, return_list=True
+    )
+
+    # TODO: copy_stack_trace
+    new_outs = [
+        xtensor_from_tensor(tensor_out, dims=output_dims) for tensor_out in tensor_outs
+    ]
+    return new_outs

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