Add a torch_tensordot operation

funsor/joint.py

@@ -214,7 +214,7 @@ def unscaled_sample(self, sampled_vars, sample_inputs):
 
 @eager.register(Joint, tuple, Funsor, Funsor)
 def eager_joint(deltas, discrete, gaussian):
-    # Demote a Joint to a simpler elementart funsor.
+    # Demote a Joint to a simpler elementary funsor.
     if not deltas:
         if gaussian is Number(0):
             return discrete

funsor/numpy.py

@@ -52,7 +52,7 @@ def align_arrays(*args):
     This is mainly useful for implementing eager funsor operations.
 
     :param funsor.terms.Funsor \*args: Multiple :class:`Array` s and
-        :class:`~funsor.terms.Number`s.
+        :class:`~funsor.terms.Number` s.
     :return: a pair ``(inputs, arrays)`` where arrayss are all
         :class:`numpy.ndarray` s that can be broadcast together to a single data
         with given ``inputs``.

funsor/torch.py

@@ -59,7 +59,7 @@ def align_tensors(*args):
     This is mainly useful for implementing eager funsor operations.
 
     :param funsor.terms.Funsor \*args: Multiple :class:`Tensor` s and
-        :class:`~funsor.terms.Number`s.
+        :class:`~funsor.terms.Number` s.
     :return: a pair ``(inputs, tensors)`` where tensors are all
         :class:`torch.Tensor` s that can be broadcast together to a single data
         with given ``inputs``.
@@ -649,6 +649,26 @@ def torch_einsum(equation, *operands):
     return Function(fn, output, operands)
 
 
+def torch_tensordot(x, y, dims):
+    """
+    Wrapper around :func:`torch.tensordot` to operate on real-valued Funsors.
+
+    Note this operates only on the ``output`` tensor. To perform sum-product
+    contractions on named dimensions, instead use ``+`` and
+    :class:`~funsor.terms.Reduce`.
+    """
+    assert isinstance(x, Funsor) and x.dtype == 'real'
+    assert isinstance(y, Funsor) and y.dtype == 'real'
+    assert isinstance(dims, int) and dims >= 0
+    x_shape = x.output.shape
+    y_shape = y.output.shape
+    assert x_shape[len(x_shape) - dims:] == y_shape[:dims]
+    shape = x_shape[:len(x_shape) - dims] + y_shape[dims:]
+    output = reals(*shape)
+    fn = functools.partial(torch.tensordot, dims=dims)
+    return Function(fn, output, (x, y))
+
+
 ################################################################################
 # Register Ops
 ################################################################################
@@ -766,4 +786,5 @@ def _safediv(x, y):
     'function',
     'materialize',
     'torch_einsum',
+    'torch_tensordot',
 ]

test/test_torch.py

@@ -11,7 +11,7 @@
 from funsor.domains import Domain, bint, find_domain, reals
 from funsor.terms import Lambda, Number, Variable
 from funsor.testing import assert_close, assert_equiv, check_funsor, random_tensor
-from funsor.torch import REDUCE_OP_TO_TORCH, Tensor, align_tensors, torch_einsum
+from funsor.torch import REDUCE_OP_TO_TORCH, Tensor, align_tensors, torch_einsum, torch_tensordot
 
 
 @pytest.mark.parametrize('shape', [(), (4,), (3, 2)])
@@ -626,3 +626,15 @@ def test_einsum(equation):
     expected = Tensor(torch.einsum(equation, *tensors))
     actual = torch_einsum(equation, *funsors)
     assert_close(actual, expected)
+
+
+@pytest.mark.parametrize('y_shape', [(), (4,), (4, 5)], ids=str)
+@pytest.mark.parametrize('xy_shape', [(), (6,), (6, 7)], ids=str)
+@pytest.mark.parametrize('x_shape', [(), (2,), (2, 3)], ids=str)
+def test_tensordot(x_shape, xy_shape, y_shape):
+    x = torch.randn(x_shape + xy_shape)
+    y = torch.randn(xy_shape + y_shape)
+    dim = len(xy_shape)
+    actual = torch_tensordot(Tensor(x), Tensor(y), dim)
+    expected = Tensor(torch.tensordot(x, y, dim))
+    assert_close(actual, expected)