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,199 @@
+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 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)
+
+
+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,72 @@
+from collections.abc import Sequence
+from typing import Literal
+
+from pytensor.tensor.slinalg import Cholesky, Solve
+from pytensor.xtensor import as_xtensor
+from pytensor.xtensor.basic import XBlockwise
+
+
+def cholesky(
+    x,
+    lower: bool = True,
+    *,
+    check_finite: bool = False,
+    overwrite_a: bool = False,
+    on_error: Literal["raise", "nan"] = "raise",
+    dims: Sequence[str],
+):
+    if len(dims) != 2:
+        raise ValueError(f"Cholesky needs two dims, got {len(dims)}")
+
+    core_op = Cholesky(
+        lower=lower,
+        check_finite=check_finite,
+        overwrite_a=overwrite_a,
+        on_error=on_error,
+    )
+    core_dims = (
+        ((dims[0], dims[1]),),
+        ((dims[0], dims[1]),),
+    )
+    x_op = XBlockwise(core_op, signature=core_op.gufunc_signature, core_dims=core_dims)
+    return x_op(x)
+
+
+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,2 @@
+import pytensor.xtensor.rewriting.basic
+import pytensor.xtensor.rewriting.shape