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,70 @@
+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, 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)
+    input_core_dims: tuple[tuple[str, str], tuple[str] | tuple[str, str]]
+    output_core_dims: tuple[tuple[str] | tuple[str, str]]
+    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,))
+        # The shared dim disappears in the output
+        output_core_dims = ((m1_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))
+        # The shared dim disappears in the output
+        output_core_dims = ((m1_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,
+        core_dims=(input_core_dims, output_core_dims),
+    )
+    return x_op(a, b)

pytensor/xtensor/math.py

@@ -0,0 +1,109 @@
+import sys
+
+import pytensor.scalar as ps
+from pytensor.scalar import ScalarOp
+from pytensor.xtensor.vectorization import XElemwise
+
+
+this_module = sys.modules[__name__]
+
+
+def _as_xelemwise(core_op: ScalarOp) -> XElemwise:
+    out = XElemwise(core_op)
+    out.__doc__ = f"Ufunc version of {core_op} for XTensorVariables"
+    return out
+
+
+abs = _as_xelemwise(ps.abs)
+add = _as_xelemwise(ps.add)
+logical_and = bitwise_and = and_ = _as_xelemwise(ps.and_)
+angle = _as_xelemwise(ps.angle)
+arccos = _as_xelemwise(ps.arccos)
+arccosh = _as_xelemwise(ps.arccosh)
+arcsin = _as_xelemwise(ps.arcsin)
+arcsinh = _as_xelemwise(ps.arcsinh)
+arctan = _as_xelemwise(ps.arctan)
+arctan2 = _as_xelemwise(ps.arctan2)
+arctanh = _as_xelemwise(ps.arctanh)
+betainc = _as_xelemwise(ps.betainc)
+betaincinv = _as_xelemwise(ps.betaincinv)
+ceil = _as_xelemwise(ps.ceil)
+clip = _as_xelemwise(ps.clip)
+complex = _as_xelemwise(ps.complex)
+conjugate = conj = _as_xelemwise(ps.conj)
+cos = _as_xelemwise(ps.cos)
+cosh = _as_xelemwise(ps.cosh)
+deg2rad = _as_xelemwise(ps.deg2rad)
+equal = eq = _as_xelemwise(ps.eq)
+erf = _as_xelemwise(ps.erf)
+erfc = _as_xelemwise(ps.erfc)
+erfcinv = _as_xelemwise(ps.erfcinv)
+erfcx = _as_xelemwise(ps.erfcx)
+erfinv = _as_xelemwise(ps.erfinv)
+exp = _as_xelemwise(ps.exp)
+exp2 = _as_xelemwise(ps.exp2)
+expm1 = _as_xelemwise(ps.expm1)
+floor = _as_xelemwise(ps.floor)
+floor_divide = floor_div = int_div = _as_xelemwise(ps.int_div)
+gamma = _as_xelemwise(ps.gamma)
+gammainc = _as_xelemwise(ps.gammainc)
+gammaincc = _as_xelemwise(ps.gammaincc)
+gammainccinv = _as_xelemwise(ps.gammainccinv)
+gammaincinv = _as_xelemwise(ps.gammaincinv)
+gammal = _as_xelemwise(ps.gammal)
+gammaln = _as_xelemwise(ps.gammaln)
+gammau = _as_xelemwise(ps.gammau)
+greater_equal = ge = _as_xelemwise(ps.ge)
+greater = gt = _as_xelemwise(ps.gt)
+hyp2f1 = _as_xelemwise(ps.hyp2f1)
+i0 = _as_xelemwise(ps.i0)
+i1 = _as_xelemwise(ps.i1)
+identity = _as_xelemwise(ps.identity)
+imag = _as_xelemwise(ps.imag)
+logical_not = bitwise_invert = bitwise_not = invert = _as_xelemwise(ps.invert)
+isinf = _as_xelemwise(ps.isinf)
+isnan = _as_xelemwise(ps.isnan)
+iv = _as_xelemwise(ps.iv)
+ive = _as_xelemwise(ps.ive)
+j0 = _as_xelemwise(ps.j0)
+j1 = _as_xelemwise(ps.j1)
+jv = _as_xelemwise(ps.jv)
+kve = _as_xelemwise(ps.kve)
+less_equal = le = _as_xelemwise(ps.le)
+log = _as_xelemwise(ps.log)
+log10 = _as_xelemwise(ps.log10)
+log1mexp = _as_xelemwise(ps.log1mexp)
+log1p = _as_xelemwise(ps.log1p)
+log2 = _as_xelemwise(ps.log2)
+less = lt = _as_xelemwise(ps.lt)
+mod = _as_xelemwise(ps.mod)
+multiply = mul = _as_xelemwise(ps.mul)
+negative = neg = _as_xelemwise(ps.neg)
+not_equal = neq = _as_xelemwise(ps.neq)
+logical_or = bitwise_or = or_ = _as_xelemwise(ps.or_)
+owens_t = _as_xelemwise(ps.owens_t)
+polygamma = _as_xelemwise(ps.polygamma)
+power = pow = _as_xelemwise(ps.pow)
+psi = _as_xelemwise(ps.psi)
+rad2deg = _as_xelemwise(ps.rad2deg)
+real = _as_xelemwise(ps.real)
+reciprocal = _as_xelemwise(ps.reciprocal)
+round = _as_xelemwise(ps.round_half_to_even)
+maximum = _as_xelemwise(ps.scalar_maximum)
+minimum = _as_xelemwise(ps.scalar_minimum)
+second = _as_xelemwise(ps.second)
+sigmoid = _as_xelemwise(ps.sigmoid)
+sign = _as_xelemwise(ps.sign)
+sin = _as_xelemwise(ps.sin)
+sinh = _as_xelemwise(ps.sinh)
+softplus = _as_xelemwise(ps.softplus)
+square = sqr = _as_xelemwise(ps.sqr)
+sqrt = _as_xelemwise(ps.sqrt)
+subtract = sub = _as_xelemwise(ps.sub)
+where = switch = _as_xelemwise(ps.switch)
+tan = _as_xelemwise(ps.tan)
+tanh = _as_xelemwise(ps.tanh)
+tri_gamma = _as_xelemwise(ps.tri_gamma)
+true_divide = true_div = _as_xelemwise(ps.true_div)
+trunc = _as_xelemwise(ps.trunc)
+logical_xor = bitwise_xor = xor = _as_xelemwise(ps.xor)

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,76 @@
+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_lower_xtensor
+from pytensor.xtensor.vectorization import XBlockwise, XElemwise
+
+
+@register_lower_xtensor
+@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_lower_xtensor
+@node_rewriter(tracks=[XBlockwise])
+def lower_blockwise(fgraph, node):
+    op: XBlockwise = node.op
+    batch_ndim = node.outputs[0].type.ndim - len(op.core_dims[1][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)
+
+    signature = op.signature or getattr(op.core_op, "gufunc_signature", None)
+    if signature is None:
+        # Build a signature based on the core dimensions
+        # The Op signature could be more strict, as core_dims will never be repeated, but no functionality depends greatly on it
+        inputs_core_dims, outputs_core_dims = op.core_dims
+        inputs_signature = ",".join(
+            f"({', '.join(inp_core_dims)})" for inp_core_dims in inputs_core_dims
+        )
+        outputs_signature = ",".join(
+            f"({', '.join(out_core_dims)})" for out_core_dims in outputs_core_dims
+        )
+        signature = f"{inputs_signature}->{outputs_signature}"
+    tensor_op = Blockwise(core_op=op.core_op, signature=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