API: Implement tensordot and matmul

Author: mtsokol

src/finch/__init__.py

@@ -15,13 +15,14 @@
     Tensor,
     astype,
     random,
+    tensordot,
+    matmul,
     permute_dims,
     multiply,
     sum,
     prod,
     add,
     subtract,
-    multiply,
     divide,
     positive,
     negative,
@@ -65,6 +66,8 @@
     "DenseStorage",
     "astype",
     "random",
+    "tensordot",
+    "matmul",
     "permute_dims",
     "int_",
     "int8",
@@ -82,7 +85,6 @@
     "complex64",
     "complex128",
     "bool",
-    "multiply",
     "lazy",
     "compiled",
     "compute",

src/finch/julia.py

@@ -1,6 +1,6 @@
 import juliapkg
 
-juliapkg.add("Finch", "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", version="0.6.19")
+juliapkg.add("Finch", "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", version="0.6.20")
 import juliacall  # noqa
 
 juliapkg.resolve()

src/finch/tensor.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional, Union
+from typing import Callable, Iterable, Optional, Union
 
 import numpy as np
 from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
@@ -113,7 +113,11 @@ def __pow__(self, other):
         return self._elemwise_op(".^", other)
 
     def __matmul__(self, other):
-        raise NotImplementedError
+        # TODO: Implement and use mul instead of tensordot
+        # https://github.com/willow-ahrens/finch-tensor/pull/22#issuecomment-2007884763
+        if self.ndim != 2 or other.ndim != 2:
+            raise ValueError(f"Both tensors must be 2-dimensional, but are: {self.ndim} and {other.ndim}.")
+        return self.tensordot(other, axes=((-1,), (-2,)))
 
     def __abs__(self):
         return self._elemwise_op("abs")
@@ -199,6 +203,15 @@ def _order(self) -> tuple[int, ...]:
     def is_computed(self) -> bool:
         return not jl.isa(self._obj, jl.Finch.LazyTensor)
 
+    def tensordot(self, other: "Tensor", /, *, axes=2) -> "Tensor":
+        if isinstance(axes, Iterable):
+            self_axes = normalize_axis_tuple(axes[0], self.ndim)
+            other_axes = normalize_axis_tuple(axes[1], other.ndim)
+            axes = (tuple(i + 1 for i in self_axes), tuple(i + 1 for i in other_axes))
+
+        result = jl.tensordot(self._obj, other._obj, axes)
+        return Tensor(result)
+
     @classmethod
     def preprocess_order(
         cls, order: OrderType, ndim: int
@@ -463,6 +476,14 @@ def prod(
     return _reduce(x, jl.prod, axis, dtype)
 
 
+def tensordot(x1: Tensor, x2: Tensor, /, *, axes=2) -> Tensor:
+    return x1.tensordot(x2, axes=axes)
+
+
+def matmul(x1: Tensor, x2: Tensor) -> Tensor:
+    return x1 @ x2
+
+
 def add(x1: Tensor, x2: Tensor, /) -> Tensor:
     return x1 + x2
 

tests/test_ops.py

@@ -86,3 +86,42 @@ def test_reductions(arr3d, func_name, axis, dtype):
         actual = actual.todense()
 
     assert_equal(actual, expected)
+
+
+def test_tensordot(arr3d):
+    A_finch = finch.Tensor(arr1d)
+    B_finch = finch.Tensor(arr2d)
+    C_finch = finch.Tensor(arr3d)
+
+    actual = finch.tensordot(B_finch, B_finch)
+    expected = np.tensordot(arr2d, arr2d)
+    assert_equal(actual.todense(), expected)
+
+    actual = finch.tensordot(B_finch, B_finch, axes=(1, 1))
+    expected = np.tensordot(arr2d, arr2d, axes=(1, 1))
+    assert_equal(actual.todense(), expected)
+
+    actual = finch.tensordot(C_finch, finch.permute_dims(C_finch, (2, 1, 0)), axes=((2, 0), (0, 2)))
+    expected = np.tensordot(arr3d, arr3d.T, axes=((2, 0), (0, 2)))
+    assert_equal(actual.todense(), expected)
+
+    actual = finch.tensordot(C_finch, A_finch, axes=(2, 0))
+    expected = np.tensordot(arr3d, arr1d, axes=(2, 0))
+    assert_equal(actual.todense(), expected)
+
+
+def test_matmul(arr2d, arr3d):
+    A_finch = finch.Tensor(arr2d)
+    B_finch = finch.Tensor(arr2d.T)
+    C_finch = finch.permute_dims(A_finch, (1, 0))
+    D_finch = finch.Tensor(arr3d)
+
+    actual = A_finch @ B_finch
+    expected = arr2d @ arr2d.T
+    assert_equal(actual.todense(), expected)
+
+    actual = A_finch @ C_finch
+    assert_equal(actual.todense(), expected)
+
+    with pytest.raises(ValueError, match="Both tensors must be 2-dimensional"):
+        A_finch @ D_finch