Generalizing XDot

Author: AllenDowney

pytensor/xtensor/math.py

@@ -162,19 +162,15 @@ def make_node(self, x, y):
         x = as_xtensor(x)
         y = as_xtensor(y)
 
-        # Filter out contracted dimensions
-        x_dims = [dim for dim in x.type.dims if dim not in self.dims]
-        y_dims = [dim for dim in y.type.dims if dim not in self.dims]
-        x_shape = [
-            size for dim, size in zip(x.type.dims, x.type.shape) if dim not in self.dims
-        ]
-        y_shape = [
-            size for dim, size in zip(y.type.dims, y.type.shape) if dim not in self.dims
-        ]
-
-        # Combine remaining dimensions
-        out_dims = tuple(x_dims + y_dims)
-        out_shape = tuple(x_shape + y_shape)
+        x_shape_dict = dict(zip(x.type.dims, x.type.shape))
+        y_shape_dict = dict(zip(y.type.dims, y.type.shape))
+        shape_dict = {**x_shape_dict, **y_shape_dict}
+
+        # Determine output dimensions
+        out_dims = tuple(d for d in shape_dict if d not in self.dims)
+
+        # Determine output shape
+        out_shape = tuple(shape_dict[d] for d in out_dims)
 
         # Determine output dtype
         out_dtype = upcast(x.type.dtype, y.type.dtype)
@@ -183,7 +179,7 @@ def make_node(self, x, y):
         return Apply(self, [x, y], [out])
 
 
-def dot(x, y, dims: str | Iterable[str] | EllipsisType | None = None):
+def dot(x, y, dim: str | Iterable[str] | EllipsisType | None = None):
     """Matrix multiplication between two XTensorVariables.
 
     This operation performs matrix multiplication between two tensors, automatically
@@ -195,7 +191,7 @@ def dot(x, y, dims: str | Iterable[str] | EllipsisType | None = None):
         First input tensor
     y : XTensorVariable
         Second input tensor
-    dims : str, Iterable[Hashable], EllipsisType, or None, optional
+    dim : str, Iterable[Hashable], EllipsisType, or None, optional
         The dimensions to contract over. If None, will contract over all matching dimensions.
         If Ellipsis (...), will contract over all dimensions.
 
@@ -214,40 +210,34 @@ def dot(x, y, dims: str | Iterable[str] | EllipsisType | None = None):
     x = as_xtensor(x)
     y = as_xtensor(y)
 
+    x_dims = set(x.type.dims)
+    y_dims = set(y.type.dims)
+    intersection = x_dims & y_dims
+    union = x_dims | y_dims
+
     # Canonicalize dims
-    if isinstance(dims, str):
-        dims = (dims,)
-    elif isinstance(dims, Iterable):
-        dims = tuple(dims)
+    if dim is None:
+        dim_set = intersection
+    elif dim is ...:
+        dim_set = union
+    elif isinstance(dim, str):
+        dim_set = {dim}
+    elif isinstance(dim, Iterable):
+        dim_set = set(dim)
 
     # Validate provided dims
-    if isinstance(dims, Iterable):
-        for dim in dims:
-            if dim not in x.type.dims:
-                raise ValueError(
-                    f"Dimension {dim} not found in first input {x.type.dims}"
-                )
-            if dim not in y.type.dims:
-                raise ValueError(
-                    f"Dimension {dim} not found in second input {y.type.dims}"
-                )
-
-    # If dims is ... , we have to sum over all remaining axes
-    sum_result = dims is ...
-
-    # Handle None and ... cases
-    if dims is None or dims is ...:
-        # Contract over all matching dimensions
-        x_dims = set(x.type.dims)
-        y_dims = set(y.type.dims)
-        dims = tuple(x_dims & y_dims)
-
-    result = XDot(dims=dims)(x, y)
-
-    if sum_result:
-        from pytensor.xtensor.reduction import sum as xtensor_sum
+    # Check if any dimension is not found in either input
+    for d in dim_set:
+        if d not in union:
+            raise ValueError(f"Dimension {d} not found in either input {y.type.dims}")
+
+    dotted_dims = tuple(dim_set & intersection)
+    summed_dims = tuple(dim_set.difference(dotted_dims))
+
+    result = XDot(dims=dotted_dims)(x, y)
 
+    if summed_dims:
         # Sum over all remaining axes
-        result = xtensor_sum(result, dim=...)
+        result = result.sum(dim=summed_dims)
 
     return result

pytensor/xtensor/type.py

@@ -650,9 +650,9 @@ def stack(self, dim, **dims):
     def unstack(self, dim, **dims):
         return px.shape.unstack(self, dim, **dims)
 
-    def dot(self, other, dims=None):
+    def dot(self, other, dim=None):
         """Matrix multiplication with another XTensorVariable, contracting over matching or specified dims."""
-        return px.math.dot(self, other, dims=dims)
+        return px.math.dot(self, other, dim=dim)
 
 
 class XTensorConstantSignature(tuple):

tests/xtensor/test_math.py

@@ -168,7 +168,7 @@ def test_dot():
     xr_assert_allclose(z_test, expected)
 
     # Test matrix-vector dot product with ellipsis
-    z = x.dot(y, dims=...)
+    z = x.dot(y, dim=...)
     fn = xr_function([x, y], z)
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim=...)
@@ -186,22 +186,22 @@ def test_dot():
     expected = x_test.dot(y_test)
     xr_assert_allclose(z_test, expected)
 
-    # Test matrix-matrix dot product with string dims
-    z = x.dot(y, dims="b")
+    # Test matrix-matrix dot product with string dim
+    z = x.dot(y, dim="b")
     fn = xr_function([x, y], z)
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim="b")
     xr_assert_allclose(z_test, expected)
 
     # Test matrix-matrix dot product with list of dims
-    z = x.dot(y, dims=["b"])
+    z = x.dot(y, dim=["b"])
     fn = xr_function([x, y], z)
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim=["b"])
     xr_assert_allclose(z_test, expected)
 
     # Test matrix-matrix dot product with ellipsis
-    z = x.dot(y, dims=...)
+    z = x.dot(y, dim=...)
     fn = xr_function([x, y], z)
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim=...)
@@ -220,28 +220,49 @@ def test_dot():
     xr_assert_allclose(z_test, expected)
 
     # Same but with explicit dimensions
-    z = x.dot(y, dims=["b", "c"])
+    z = x.dot(y, dim=["b", "c"])
     fn = xr_function([x, y], z)
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim=["b", "c"])
     xr_assert_allclose(z_test, expected)
 
     # Same but with ellipses
-    z = x.dot(y, dims=...)
+    z = x.dot(y, dim=...)
     fn = xr_function([x, y], z)
 
     z_test = fn(x_test, y_test)
     expected = x_test.dot(y_test, dim=...)
     xr_assert_allclose(z_test, expected)
 
+    # Dot product with sum
+    x_test = DataArray(np.arange(24.0).reshape(2, 3, 4), dims=("a", "b", "c"))
+    y_test = DataArray(np.arange(60.0).reshape(3, 4, 5), dims=("b", "c", "d"))
+    expected = x_test.dot(y_test, dim=("a", "b", "c"))
+
+    x = xtensor("x", dims=("a", "b", "c"), shape=(2, 3, 4))
+    y = xtensor("y", dims=("b", "c", "d"), shape=(3, 4, 5))
+    z = x.dot(y, dim=("a", "b", "c"))
+    fn = xr_function([x, y], z)
+    z_test = fn(x_test, y_test)
+    xr_assert_allclose(z_test, expected)
+
+    return
+    # Dot product with sum in the middle
+    # This is not supported yet
+    x_test = DataArray(np.arange(120).reshape(2, 3, 4, 5), dims=("a", "b", "c", "d"))
+    y_test = DataArray(np.arange(360).reshape(3, 4, 5, 6), dims=("b", "c", "d", "e"))
+    expected = x_test.dot(y_test, dim=("b", "d"))
+    x = xtensor("x", dims=("a", "b", "c", "d"), shape=(2, 3, 4, 5))
+    y = xtensor("y", dims=("b", "c", "d", "e"), shape=(3, 4, 5, 6))
+    z = x.dot(y, dim=("b", "d"))
+    fn = xr_function([x, y], z)
+    z_test = fn(x_test, y_test)
+    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"))