Implement Elemwise and Blockwise operations for XTensorVariables

Author: ricardoV94

pytensor/xtensor/__init__.py

@@ -1,6 +1,9 @@
 import warnings
 
 import pytensor.xtensor.rewriting
+from pytensor.xtensor import (
+    linalg,
+)
 from pytensor.xtensor.type import (
     XTensorType,
     as_xtensor,

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.type import as_xtensor
+from pytensor.xtensor.vectorization 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.vectorization 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

@@ -1,2 +1,3 @@
 import pytensor.xtensor.rewriting.basic
 import pytensor.xtensor.rewriting.shape
+import pytensor.xtensor.rewriting.vectorization

pytensor/xtensor/rewriting/vectorization.py

@@ -0,0 +1,64 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.xtensor.basic import tensor_from_xtensor, xtensor_from_tensor
+from pytensor.xtensor.rewriting.utils import register_xcanonicalize
+from pytensor.xtensor.vectorization import XBlockwise, XElemwise
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[XElemwise])
+def lower_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 lower_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/type.py

@@ -146,6 +146,109 @@ def __complex__(self):
             "Call `.astype(complex)` for the symbolic equivalent."
         )
 
+    # Python valid overloads
+    def __abs__(self):
+        return px.math.abs(self)
+
+    def __neg__(self):
+        return px.math.neg(self)
+
+    def __lt__(self, other):
+        return px.math.lt(self, other)
+
+    def __le__(self, other):
+        return px.math.le(self, other)
+
+    def __gt__(self, other):
+        return px.math.gt(self, other)
+
+    def __ge__(self, other):
+        return px.math.ge(self, other)
+
+    def __invert__(self):
+        return px.math.invert(self)
+
+    def __and__(self, other):
+        return px.math.and_(self, other)
+
+    def __or__(self, other):
+        return px.math.or_(self, other)
+
+    def __xor__(self, other):
+        return px.math.xor(self, other)
+
+    def __rand__(self, other):
+        return px.math.and_(other, self)
+
+    def __ror__(self, other):
+        return px.math.or_(other, self)
+
+    def __rxor__(self, other):
+        return px.math.xor(other, self)
+
+    def __add__(self, other):
+        return px.math.add(self, other)
+
+    def __sub__(self, other):
+        return px.math.sub(self, other)
+
+    def __mul__(self, other):
+        return px.math.mul(self, other)
+
+    def __div__(self, other):
+        return px.math.div(self, other)
+
+    def __pow__(self, other):
+        return px.math.pow(self, other)
+
+    def __mod__(self, other):
+        return px.math.mod(self, other)
+
+    def __divmod__(self, other):
+        return px.math.divmod(self, other)
+
+    def __truediv__(self, other):
+        return px.math.true_div(self, other)
+
+    def __floordiv__(self, other):
+        return px.math.floor_div(self, other)
+
+    def __rtruediv__(self, other):
+        return px.math.true_div(other, self)
+
+    def __rfloordiv__(self, other):
+        return px.math.floor_div(other, self)
+
+    def __radd__(self, other):
+        return px.math.add(other, self)
+
+    def __rsub__(self, other):
+        return px.math.sub(other, self)
+
+    def __rmul__(self, other):
+        return px.math.mul(other, self)
+
+    def __rdiv__(self, other):
+        return px.math.div_proxy(other, self)
+
+    def __rmod__(self, other):
+        return px.math.mod(other, self)
+
+    def __rdivmod__(self, other):
+        return px.math.divmod(other, self)
+
+    def __rpow__(self, other):
+        return px.math.pow(other, self)
+
+    def __ceil__(self):
+        return px.math.ceil(self)
+
+    def __floor__(self):
+        return px.math.floor(self)
+
+    def __trunc__(self):
+        return px.math.trunc(self)
+
     # DataArray-like attributes
     # https://docs.xarray.dev/en/latest/api.html#id1
     @property
@@ -192,14 +295,33 @@ def dtype(self):
 
     # DataArray contents
     # https://docs.xarray.dev/en/latest/api.html#dataarray-contents
-    def rename(self, new_name_or_name_dict, **names):
+    def rename(self, new_name_or_name_dict=None, **names):
         if isinstance(new_name_or_name_dict, str):
-            # TODO: Should we make a symbolic copy?
-            self.name = new_name_or_name_dict
+            new_name = new_name_or_name_dict
             name_dict = None
         else:
+            new_name = None
             name_dict = new_name_or_name_dict
-        return px.basic.rename(name_dict, **names)
+        new_out = px.basic.rename(self, name_dict, **names)
+        new_out.name = new_name
+        return new_out
+
+    # def swap_dims(self, *args, **kwargs):
+    #     ...
+    #
+    # def expand_dims(self, *args, **kwargs):
+    #     ...
+    #
+    # def squeeze(self):
+    #     ...
+
+    def copy(self, name: str | None = None):
+        out = px.math.identity(self)
+        out.name = name
+        return out
+
+    def astype(self, dtype):
+        return px.math.cast(self, dtype)
 
     def item(self):
         raise NotImplementedError("item not implemented for XTensorVariable")
@@ -219,6 +341,26 @@ def sel(self, *args, **kwargs):
     def __getitem__(self, idx):
         raise NotImplementedError("Indexing not yet implemnented")
 
+    # ndarray methods
+    # https://docs.xarray.dev/en/latest/api.html#id7
+    def clip(self, min, max):
+        return px.math.clip(self, min, max)
+
+    def conj(self):
+        return px.math.conj(self)
+
+    @property
+    def imag(self):
+        return px.math.imag(self)
+
+    @property
+    def real(self):
+        return px.math.real(self)
+
+    # @property
+    # def T(self):
+    #     ...
+
 
 class XTensorConstantSignature(tuple):
     def __eq__(self, other):