Adding shape checking at rewrite time

Author: AllenDowney

pytensor/xtensor/math.py

@@ -1,5 +1,5 @@
 import sys
-from collections.abc import Hashable, Iterable
+from collections.abc import Iterable
 from types import EllipsisType
 
 import numpy as np
@@ -190,7 +190,7 @@ def make_node(self, x, y):
         return Apply(self, [x, y], [out])
 
 
-def dot(x, y, dims: str | Iterable[Hashable] | EllipsisType | None = None):
+def dot(x, y, dims: str | Iterable[str] | EllipsisType | None = None):
     """Matrix multiplication between two XTensorVariables.
 
     This operation performs matrix multiplication between two tensors, automatically

pytensor/xtensor/rewriting/math.py

@@ -20,11 +20,21 @@ def lower_dot(fgraph, node):
     x_tensor = tensor_from_xtensor(x)
     y_tensor = tensor_from_xtensor(y)
 
-    # Get axes to contract for each input
+    # Get the axes for contraction
     x_axes = [x.type.dims.index(dim) for dim in node.op.dims]
     y_axes = [y.type.dims.index(dim) for dim in node.op.dims]
 
-    # Perform dot product
+    # Check that shapes match along contracted dimensions
+    for dim in node.op.dims:
+        x_idx = x.type.dims.index(dim)
+        y_idx = y.type.dims.index(dim)
+        if x.type.shape[x_idx] != y.type.shape[y_idx]:
+            raise ValueError(
+                "Input arrays have inconsistent type shape along the axes "
+                f"that are to be reduced with tensordot: {x.type.shape[x_idx]} != {y.type.shape[y_idx]}"
+            )
+
+    # Perform the tensordot operation
     out_tensor = tensordot(x_tensor, y_tensor, axes=(x_axes, y_axes))
 
     # Sum over all remaining axes if needed

tests/xtensor/test_math.py

@@ -229,3 +229,29 @@ def test_dot():
         z_test = fn(x_test, y_test)
         expected = x_test.dot(y_test, dim=...)
         xr_assert_allclose(z_test, expected)
+
+
+def test_dot_errors():
+    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
+    y = xtensor("y", dims=("b", "c"), shape=(3, 4))
+    with pytest.raises(ValueError, match="Dimension c not found in first input"):
+        x.dot(y, dims=["c"])
+    with pytest.raises(ValueError, match="Dimension a not found in second input"):
+        x.dot(y, dims=["a"])
+
+    # Test a case where there are no matching dimensions
+    x_test = DataArray(np.ones((2, 3)), dims=("a", "b"))
+    y_test = DataArray(np.ones((4, 5)), dims=("b", "c"))
+    with pytest.raises(ValueError, match="cannot reindex or align along dimension"):
+        x_test.dot(y_test)
+
+    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
+    y = xtensor("y", dims=("b", "c"), shape=(4, 5))
+    with pytest.raises(
+        ValueError, match="Input arrays have inconsistent type shape along the axes"
+    ):
+        z = x.dot(y)
+        fn = function([x, y], z)
+        x_test = DataArray(np.ones((2, 3)), dims=("a", "b"))
+        y_test = DataArray(np.ones((4, 5)), dims=("b", "c"))
+        fn(x_test, y_test)