Refactor XElemwise and XBlockwise

Author: ricardoV94

pytensor/xtensor/linalg.py

@@ -28,7 +28,7 @@ def cholesky(
         ((dims[0], dims[1]),),
         ((dims[0], dims[1]),),
     )
-    x_op = XBlockwise(core_op, signature=core_op.gufunc_signature, core_dims=core_dims)
+    x_op = XBlockwise(core_op, core_dims=core_dims)
     return x_op(x)
 
 
@@ -48,18 +48,15 @@ def solve(
         [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,),)
+        # 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))
-        output_core_dims = (
-            (
-                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")
 
@@ -68,7 +65,6 @@ def solve(
     )
     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/rewriting/vectorization.py

@@ -37,7 +37,7 @@ def lower_elemwise(fgraph, node):
 @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_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
@@ -53,7 +53,19 @@ def lower_blockwise(fgraph, node):
         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)
+    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

pytensor/xtensor/vectorization.py

@@ -3,11 +3,27 @@
 from pytensor import scalar as ps
 from pytensor.graph import Apply, Op
 from pytensor.tensor import tensor
-from pytensor.tensor.utils import _parse_gufunc_signature
 from pytensor.xtensor.basic import XOp
 from pytensor.xtensor.type import as_xtensor, xtensor
 
 
+def combine_dims_and_shape(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
+    return dims_and_shape
+
+
 class XElemwise(XOp):
     __props__ = ("scalar_op",)
 
@@ -22,20 +38,7 @@ def make_node(self, *inputs):
                 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
-
+        dims_and_shape = combine_dims_and_shape(inputs)
         if dims_and_shape:
             output_dims, output_shape = zip(*dims_and_shape.items())
         else:
@@ -53,48 +56,33 @@ def make_node(self, *inputs):
 
 
 class XBlockwise(XOp):
-    __props__ = ("core_op", "signature", "core_dims")
+    __props__ = ("core_op", "core_dims")
 
     def __init__(
         self,
         core_op: Op,
-        signature: str,
         core_dims: tuple[tuple[tuple[str, ...], ...], tuple[tuple[str, ...], ...]],
+        signature: str | None = None,
     ):
         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
+        self.signature = signature  # Only used for lowering, not for validation
 
     def make_node(self, *inputs):
         inputs = [as_xtensor(i) for i in inputs]
-        if len(inputs) != len(self.inputs_sig):
+        if len(inputs) != len(self.core_dims[0]):
             raise ValueError(
-                f"Wrong number of inputs, expected {len(self.inputs_sig)}, got {len(inputs)}"
+                f"Wrong number of inputs, expected {len(self.core_dims[0])}, 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
+        dims_and_shape = combine_dims_and_shape(inputs)
 
         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())
+        core_input_dims_set = set(chain.from_iterable(core_inputs_dims))
+        batch_dims, batch_shape = zip(
+            *((k, v) for k, v in dims_and_shape.items() if k not in core_input_dims_set)
+        )
 
         dummy_core_inputs = []
         for inp, core_inp_dims in zip(inputs, core_inputs_dims):

tests/xtensor/test_linalg.py

@@ -1,5 +1,6 @@
 # ruff: noqa: E402
 import pytest
+from xtensor.util import xr_assert_allclose, xr_function
 
 
 pytest.importorskip("xarray")
@@ -14,7 +15,6 @@
     solve as xr_solve,
 )
 
-from pytensor import function
 from pytensor.xtensor.linalg import cholesky, solve
 from pytensor.xtensor.type import xtensor
 
@@ -25,59 +25,52 @@ def test_cholesky():
     assert y.type.dims == ("batch", "b", "a")
     assert y.type.shape == (3, 4, 4)
 
-    fn = function([x], y)
+    fn = xr_function([x], y)
     rng = np.random.default_rng(25)
-    x_ = rng.random(size=(4, 3, 3))
+    x_ = rng.random(size=(3, 4, 4))
     x_ = x_ @ x_.mT
     x_test = DataArray(x_.transpose(1, 0, 2), dims=x.type.dims)
-    np.testing.assert_allclose(
-        fn(x_test.values),
-        xr_cholesky(x_test, dims=["b", "a"]).values,
+    xr_assert_allclose(
+        fn(x_test),
+        xr_cholesky(x_test, dims=["b", "a"]),
     )
 
 
 def test_solve_vector_b():
     a = xtensor("a", dims=("city", "country", "galaxy"), shape=(None, 4, 1))
     b = xtensor("b", dims=("city", "planet"), shape=(None, 2))
     x = solve(a, b, dims=["country", "city"])
-    assert x.type.dims == ("galaxy", "planet", "city")
-    assert x.type.shape == (
-        1,
-        2,
-        None,
-    )  # Core Solve doesn't make use of the fact A must be square in the static shape
+    assert x.type.dims == ("galaxy", "planet", "country")
+    # Core Solve doesn't make use of the fact A must be square in the static shape
+    assert x.type.shape == (1, 2, None)
 
-    fn = function([a, b], x)
+    fn = xr_function([a, b], x)
 
     rng = np.random.default_rng(25)
     a_test = DataArray(rng.random(size=(4, 4, 1)), dims=a.type.dims)
     b_test = DataArray(rng.random(size=(4, 2)), dims=b.type.dims)
 
-    np.testing.assert_allclose(
-        fn(a_test.values, b_test.values),
-        xr_solve(a_test, b_test, dims=["country", "city"]).values,
+    xr_assert_allclose(
+        fn(a_test, b_test),
+        xr_solve(a_test, b_test, dims=["country", "city"]),
     )
 
 
 def test_solve_matrix_b():
     a = xtensor("a", dims=("city", "country", "galaxy"), shape=(None, 4, 1))
     b = xtensor("b", dims=("district", "city", "planet"), shape=(5, None, 2))
     x = solve(a, b, dims=["country", "city", "district"])
-    assert x.type.dims == ("galaxy", "planet", "city", "district")
-    assert x.type.shape == (
-        1,
-        2,
-        None,
-        5,
-    )  # Core Solve doesn't make use of the fact A must be square in the static shape
+    assert x.type.dims == ("galaxy", "planet", "country", "district")
+    # Core Solve doesn't make use of the fact A must be square in the static shape
+    assert x.type.shape == (1, 2, None, 5)
 
-    fn = function([a, b], x)
+    fn = xr_function([a, b], x)
 
     rng = np.random.default_rng(25)
     a_test = DataArray(rng.random(size=(4, 4, 1)), dims=a.type.dims)
     b_test = DataArray(rng.random(size=(5, 4, 2)), dims=b.type.dims)
 
-    np.testing.assert_allclose(
-        fn(a_test.values, b_test.values),
-        xr_solve(a_test, b_test, dims=["country", "city", "district"]).values,
+    xr_assert_allclose(
+        fn(a_test, b_test),
+        xr_solve(a_test, b_test, dims=["country", "city", "district"]),
     )