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