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,85 @@
+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):
+        inputs = [px.as_xtensor(inp) for inp in inputs]
+
+        # TODO: This ordering is different than what xarray does
+        unique_dims: dict[str, int | None] = {}
+        for inp in inputs:
+            for dim, dim_length in zip(inp.type.dims, inp.type.shape):
+                if dim not in unique_dims:
+                    unique_dims[dim] = dim_length
+                elif dim_length is not None:
+                    # Check for conflicting shapes
+                    if (unique_dims[dim] is not None) and (
+                        unique_dims[dim] != dim_length
+                    ):
+                        raise ValueError(f"Dimension {dim} has conflicting shapes")
+                    # Keep the non-None shape
+                    unique_dims[dim] = dim_length
+
+        dims, shape = zip(*sorted(unique_dims.items()))
+
+        # TODO: Fix dtype
+        output_type = px.XTensorType("float64", dims=dims, shape=shape)
+        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)
+exp = XElemwise(ps.exp)

pytensor/xtensor/rewriting/__init__.py

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

pytensor/xtensor/rewriting/basic.py

@@ -0,0 +1,54 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor import expand_dims
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.xtensor.basic import (
+    TensorFromXTensor,
+    XElemwise,
+    XTensorFromTensor,
+    tensor_from_xtensor,
+    xtensor_from_tensor,
+)
+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=[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

pytensor/xtensor/type.py

@@ -0,0 +1,250 @@
+try:
+    import xarray as xr
+
+    XARRAY_AVAILABLE = True
+except ModuleNotFoundError:
+    XARRAY_AVAILABLE = False
+
+from collections.abc import Sequence
+from typing import TypeVar
+
+import numpy as np
+
+from pytensor import _as_symbolic, config
+from pytensor import scalar as aes
+from pytensor.graph import Apply, Constant
+from pytensor.graph.basic import Variable
+from pytensor.graph.type import HasDataType, HasShape, Type
+from pytensor.tensor.utils import hash_from_ndarray
+from pytensor.utils import hash_from_code
+
+
+class XTensorType(Type, HasDataType, HasShape):
+    """A `Type` for Xtensors (Xarray-like tensors with dims)."""
+
+    __props__ = ("dtype", "shape", "dims")
+
+    def __init__(
+        self,
+        dtype: str | np.dtype,
+        *,
+        dims: Sequence[str],
+        shape: Sequence[int | None] | None = None,
+        name: str | None = None,
+    ):
+        if dtype == "floatX":
+            self.dtype = config.floatX
+        else:
+            if np.obj2sctype(dtype) is None:
+                raise TypeError(f"Invalid dtype: {dtype}")
+
+            self.dtype = np.dtype(dtype).name
+
+        self.dims = tuple(dims)
+        if shape is None:
+            self.shape = (None,) * len(self.dims)
+        else:
+            self.shape = tuple(shape)
+        self.name = name
+
+    def clone(
+        self,
+        dtype=None,
+        dims=None,
+        shape=None,
+        **kwargs,
+    ):
+        if dtype is None:
+            dtype = self.dtype
+        if dims is None:
+            dims = self.dims
+        if shape is None:
+            shape = self.shape
+        return type(self)(format, dtype, shape=shape, dims=dims, **kwargs)
+
+    def filter(self, value, strict=False, allow_downcast=None):
+        # TODO: Implement this
+        return value
+
+        if isinstance(value, Variable):
+            raise TypeError(
+                "Expected an array-like object, but found a Variable: "
+                "maybe you are trying to call a function on a (possibly "
+                "shared) variable instead of a numeric array?"
+            )
+
+        if (
+            isinstance(value, self.format_cls[self.format])
+            and value.dtype == self.dtype
+        ):
+            return value
+
+        if strict:
+            raise TypeError(
+                f"{value} is not sparse, or not the right dtype (is {value.dtype}, "
+                f"expected {self.dtype})"
+            )
+
+        # The input format could be converted here
+        if allow_downcast:
+            sp = self.format_cls[self.format](value, dtype=self.dtype)
+        else:
+            data = self.format_cls[self.format](value)
+            up_dtype = aes.upcast(self.dtype, data.dtype)
+            if up_dtype != self.dtype:
+                raise TypeError(f"Expected {self.dtype} dtype but got {data.dtype}")
+            sp = data.astype(up_dtype)
+
+        assert sp.format == self.format
+
+        return sp
+
+    def convert_variable(self, var):
+        # TODO: Implement this
+        return var
+        res = super().convert_variable(var)
+
+        if res is None:
+            return res
+
+        if not isinstance(res.type, type(self)):
+            return None
+
+        if res.dims != self.dims:
+            # TODO: Does this make sense?
+            return None
+
+        return res
+
+    def __hash__(self):
+        return hash(super().__hash__(), self.shape, self.dims)
+
+    def __repr__(self):
+        # TODO: Add `?` for unknown shapes like `TensorType` does
+        return f"XTensorType({self.dtype}, {self.dims}, {self.shape})"
+
+    def __eq__(self, other):
+        res = super().__eq__(other)
+
+        if isinstance(res, bool):
+            return res and self.dims == other.dims and self.shape == other.shape
+
+        return res
+
+    def is_super(self, otype):
+        # TODO: Implement this
+        return True
+
+        if not super().is_super(otype):
+            return False
+
+        if self.dims == otype.dims:
+            return True
+
+        return False
+
+
+def xtensor(
+    name: str | None = None,
+    *,
+    dims: Sequence[str],
+    shape: Sequence[int | None] | None = None,
+    dtype: str | np.dtype = "floatX",
+):
+    return XTensorType(dtype, dims=dims, shape=shape)(name=name)
+
+
+# class _x_tensor_py_operators
+
+
+class XTensorVariable(Variable):
+    pass
+
+    # def __str__(self):
+    #     return f"{self.__class__.__name__}{{{self.format},{self.dtype}}}"
+
+    # def __repr__(self):
+    #     return str(self)
+
+
+class XTensorConstantSignature(tuple):
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        (t0, d0), (t1, d1) = self, other
+        if t0 != t1 or d0.shape != d1.shape:
+            return False
+
+        return True
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __hash__(self):
+        (a, b) = self
+        return hash(type(self)) ^ hash(a) ^ hash(type(b))
+
+    def pytensor_hash(self):
+        t, d = self
+        return "".join([hash_from_ndarray(d)] + [hash_from_code(dim) for dim in t.dims])
+
+
+_XTensorTypeType = TypeVar("_XTensorTypeType", bound=XTensorType)
+
+
+class XTensorConstant(XTensorVariable, Constant[_XTensorTypeType]):
+    def __init__(self, type: _XTensorTypeType, data, name=None):
+        # TODO: Add checks that type and data are compatible
+        Constant.__init__(self, type, data, name)
+
+    def signature(self):
+        assert self.data is not None
+        return XTensorConstantSignature((self.type, self.data))
+
+
+XTensorType.variable_type = XTensorVariable
+XTensorType.constant_type = XTensorConstant
+
+
+def xtensor_constant(x, name=None):
+    if not isinstance(x, xr.DataArray):
+        raise TypeError("xtensor.constant must be called on a Xarray DataArray")
+    try:
+        return XTensorConstant(
+            XTensorType(dtype=x.dtype, dims=x.dims, shape=x.shape),
+            x.values.copy(),
+            name=name,
+        )
+    except TypeError:
+        raise TypeError(f"Could not convert {x} to XTensorType")
+
+
+if XARRAY_AVAILABLE:
+
+    @_as_symbolic.register(xr.DataArray)
+    def as_symbolic_xarray(x, **kwargs):
+        return xtensor_constant(x, **kwargs)
+
+
+def as_xtensor_variable(x, name=None):
+    if isinstance(x, Apply):
+        if len(x.outputs) != 1:
+            raise ValueError(
+                "It is ambiguous which output of a "
+                "multi-output Op has to be fetched.",
+                x,
+            )
+        else:
+            x = x.outputs[0]
+    if isinstance(x, Variable):
+        if not isinstance(x.type, XTensorType):
+            raise TypeError(f"Variable type field must be a XTensorType, got {x.type}")
+        return x
+    try:
+        return xtensor_constant(x, name=name)
+    except TypeError as err:
+        raise TypeError(f"Cannot convert {x} to XTensorType {type(x)}") from err
+
+
+as_xtensor = as_xtensor_variable

tests/xtensor/__init__.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from pytensor import function
+from pytensor.xtensor.basic import add, exp
+from pytensor.xtensor.type import xtensor
+
+
+def test_add():
+    x = xtensor("x", dims=("city",), shape=(None,))
+    y = xtensor("y", dims=("country",), shape=(4,))
+    z = add(exp(x), exp(y))
+    assert z.type.dims == ("city", "country")
+    assert z.type.shape == (None, 4)
+
+    fn = function([x, y], z)
+    # fn.dprint(print_type=True)
+
+    np.testing.assert_allclose(
+        fn(x=np.zeros(3), y=np.zeros(4)),
+        np.full((3, 4), 2.0),
+    )