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,167 @@
+from itertools import chain
+
+import pytensor.scalar as ps
+from pytensor.graph import Apply, Op
+from pytensor.tensor import TensorType, tensor
+from pytensor.tensor.utils import _parse_gufunc_signature
+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(
+            "xtensor operations must be rewritten as tensor operations"
+        )
+
+
+class TensorFromXTensor(Op):
+    view_map = {0: [0]}
+
+    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])
+
+    def perform(self, node, inputs, output_storage) -> None:
+        [x] = inputs
+        output_storage[0][0] = x
+
+
+tensor_from_xtensor = TensorFromXTensor()
+
+
+class XTensorFromTensor(Op):
+    view_map = {0: [0]}
+    __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 perform(self, node, inputs, output_storage) -> None:
+        [x] = inputs
+        output_storage[0][0] = x
+
+
+def xtensor_from_tensor(x, dims):
+    return XTensorFromTensor(dims=dims)(x)
+
+
+class XElemwise(XOp):
+    __props__ = ("scalar_op",)
+
+    def __init__(self, scalar_op):
+        super().__init__()
+        self.scalar_op = scalar_op
+
+    def make_node(self, *inputs):
+        inputs = [as_xtensor(inp) for inp in inputs]
+        if (self.scalar_op.nin != -1) and (len(inputs) != self.scalar_op.nin):
+            raise ValueError(
+                f"Wrong number of inputs, expected {self.scalar_op.nin}, got {len(inputs)}"
+            )
+
+        dims_and_shape: dict[str, int | None] = {}
+        for inp in inputs:
+            for dim, dim_length in zip(inp.type.dims, inp.type.shape):
+                if dim not in dims_and_shape:
+                    dims_and_shape[dim] = dim_length
+                elif dim_length is not None:
+                    # Check for conflicting shapes
+                    if (dims_and_shape[dim] is not None) and (
+                        dims_and_shape[dim] != dim_length
+                    ):
+                        raise ValueError(f"Dimension {dim} has conflicting shapes")
+                    # Keep the non-None shape
+                    dims_and_shape[dim] = dim_length
+
+        output_dims, output_shape = zip(*dims_and_shape.items())
+
+        dummy_scalars = [ps.get_scalar_type(inp.type.dtype)() for inp in inputs]
+        output_dtypes = [
+            out.type.dtype for out in self.scalar_op.make_node(*dummy_scalars).outputs
+        ]
+        outputs = [
+            xtensor(dtype=output_dtype, dims=output_dims, shape=output_shape)
+            for output_dtype in output_dtypes
+        ]
+        return Apply(self, inputs, outputs)
+
+
+class XBlockwise(XOp):
+    __props__ = ("core_op", "signature", "core_dims")
+
+    def __init__(
+        self,
+        core_op: Op,
+        signature: str,
+        core_dims: tuple[tuple[tuple[str, ...], ...], tuple[tuple[str, ...], ...]],
+    ):
+        super().__init__()
+        self.core_op = core_op
+        self.signature = signature
+        self.inputs_sig, self.outputs_sig = _parse_gufunc_signature(signature)
+        self.core_dims = core_dims
+
+    def make_node(self, *inputs):
+        inputs = [as_xtensor(i) for i in inputs]
+        if len(inputs) != len(self.inputs_sig):
+            raise ValueError(
+                f"Wrong number of inputs, expected {len(self.inputs_sig)}, got {len(inputs)}"
+            )
+
+        dims_and_shape: dict[str, int | None] = {}
+        for inp in inputs:
+            for dim, dim_length in zip(inp.type.dims, inp.type.shape):
+                if dim not in dims_and_shape:
+                    dims_and_shape[dim] = dim_length
+                elif dim_length is not None:
+                    # Check for conflicting shapes
+                    if (dims_and_shape[dim] is not None) and (
+                        dims_and_shape[dim] != dim_length
+                    ):
+                        raise ValueError(f"Dimension {dim} has conflicting shapes")
+                    # Keep the non-None shape
+                    dims_and_shape[dim] = dim_length
+
+        core_inputs_dims, core_outputs_dims = self.core_dims
+        # TODO: Avoid intermediate dict
+        core_dims = set(chain.from_iterable(core_inputs_dims))
+        batched_dims_and_shape = {
+            k: v for k, v in dims_and_shape.items() if k not in core_dims
+        }
+        batch_dims, batch_shape = zip(*batched_dims_and_shape.items())
+
+        dummy_core_inputs = []
+        for inp, core_inp_dims in zip(inputs, core_inputs_dims):
+            try:
+                core_static_shape = [
+                    inp.type.shape[inp.type.dims.index(d)] for d in core_inp_dims
+                ]
+            except IndexError:
+                raise ValueError(
+                    f"At least one core dim={core_inp_dims} missing from input {inp} with dims={inp.type.dims}"
+                )
+            dummy_core_inputs.append(
+                tensor(dtype=inp.type.dtype, shape=core_static_shape)
+            )
+        core_node = self.core_op.make_node(*dummy_core_inputs)
+
+        outputs = [
+            xtensor(
+                dtype=core_out.type.dtype,
+                shape=batch_shape + core_out.type.shape,
+                dims=batch_dims + core_out_dims,
+            )
+            for core_out, core_out_dims in zip(core_node.outputs, core_outputs_dims)
+        ]
+        return Apply(self, inputs, outputs)

pytensor/xtensor/linalg.py

@@ -0,0 +1,45 @@
+from collections.abc import Sequence
+
+from pytensor.tensor.slinalg import Solve
+from pytensor.xtensor import as_xtensor
+from pytensor.xtensor.basic import XBlockwise
+
+
+def solve(
+    a,
+    b,
+    dims: Sequence[str],
+    assume_a="gen",
+    lower: bool = False,
+    check_finite: bool = False,
+):
+    a, b = as_xtensor(a), as_xtensor(b)
+    if len(dims) == 2:
+        b_ndim = 1
+        [m1_dim] = [dim for dim in dims if dim not in b.type.dims]
+        m2_dim = dims[0] if dims[0] != m1_dim else dims[1]
+        input_core_dims = ((m1_dim, m2_dim), (m2_dim,))
+        output_core_dims = ((m2_dim,),)
+    elif len(dims) == 3:
+        b_ndim = 2
+        [n_dim] = [dim for dim in dims if dim not in a.type.dims]
+        [m1_dim, m2_dim] = [dim for dim in dims if dim != n_dim]
+        input_core_dims = ((m1_dim, m2_dim), (m2_dim, n_dim))
+        output_core_dims = (
+            (
+                m2_dim,
+                n_dim,
+            ),
+        )
+    else:
+        raise ValueError("Solve dims must have length 2 or 3")
+
+    core_op = Solve(
+        b_ndim=b_ndim, assume_a=assume_a, lower=lower, check_finite=check_finite
+    )
+    x_op = XBlockwise(
+        core_op,
+        signature=core_op.gufunc_signature,
+        core_dims=(input_core_dims, output_core_dims),
+    )
+    return x_op(a, b)

pytensor/xtensor/math.py

@@ -0,0 +1,31 @@
+import inspect
+import sys
+
+import pytensor.scalar as ps
+from pytensor.scalar import ScalarOp
+from pytensor.xtensor.basic import XElemwise
+
+
+this_module = sys.modules[__name__]
+
+
+def get_all_scalar_ops():
+    """
+    Find all scalar operations in the pytensor.scalar module that can be wrapped with XElemwise.
+
+    Returns:
+        dict: A dictionary mapping operation names to XElemwise instances
+    """
+    result = {}
+
+    # Get all module members
+    for name, obj in inspect.getmembers(ps):
+        # Check if the object is a scalar op (has make_node method and is not an abstract class)
+        if isinstance(obj, ScalarOp):
+            result[name] = XElemwise(obj)
+
+    return result
+
+
+for name, op in get_all_scalar_ops().items():
+    setattr(this_module, name, op)

pytensor/xtensor/rewriting/__init__.py

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

pytensor/xtensor/rewriting/basic.py

@@ -0,0 +1,88 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.xtensor.basic import (
+    TensorFromXTensor,
+    XBlockwise,
+    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):
+    out_dims = node.outputs[0].type.dims
+
+    # Convert input XTensors to Tensors and align batch dimensions
+    tensor_inputs = []
+    for inp in node.inputs:
+        inp_dims = inp.type.dims
+        order = [
+            inp_dims.index(out_dim) if out_dim in inp_dims else "x"
+            for out_dim in out_dims
+        ]
+        tensor_inp = tensor_from_xtensor(inp).dimshuffle(order)
+        tensor_inputs.append(tensor_inp)
+
+    tensor_outs = Elemwise(scalar_op=node.op.scalar_op)(
+        *tensor_inputs, return_list=True
+    )
+
+    # Convert output Tensors to XTensors
+    new_outs = [
+        xtensor_from_tensor(tensor_out, dims=out_dims) for tensor_out in tensor_outs
+    ]
+    return new_outs
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[XBlockwise])
+def xblockwise_to_blockwise(fgraph, node):
+    op: XBlockwise = node.op
+    batch_ndim = node.outputs[0].type.ndim - len(op.outputs_sig[0])
+    batch_dims = node.outputs[0].type.dims[:batch_ndim]
+
+    # Convert input Tensors to XTensors, align batch dimensions and place core dimension at the end
+    tensor_inputs = []
+    for inp, core_dims in zip(node.inputs, op.core_dims[0]):
+        inp_dims = inp.type.dims
+        # Align the batch dims of the input, and place the core dims on the right
+        batch_order = [
+            inp_dims.index(batch_dim) if batch_dim in inp_dims else "x"
+            for batch_dim in batch_dims
+        ]
+        core_order = [inp_dims.index(core_dim) for core_dim in core_dims]
+        tensor_inp = tensor_from_xtensor(inp).dimshuffle(batch_order + core_order)
+        tensor_inputs.append(tensor_inp)
+
+    tensor_op = Blockwise(core_op=node.op.core_op, signature=op.signature)
+    tensor_outs = tensor_op(*tensor_inputs, return_list=True)
+
+    # Convert output Tensors to XTensors
+    new_outs = [
+        xtensor_from_tensor(tensor_out, dims=old_out.type.dims)
+        for (tensor_out, old_out) in zip(tensor_outs, node.outputs, strict=True)
+    ]
+    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