Add oneflow.tensordot interface

docs/source/oneflow.rst

@@ -148,6 +148,7 @@ oneflow
             tan, 
             tanh, 
             tensor, 
+            tensordot,
             tile, 
             transpose,
             t,

oneflow/core/functional/functional_api.yaml

@@ -902,6 +902,11 @@
     Double alpha=1.0) => BatchMatMul"
   bind_python: True
 
+- name: "tensordot"
+  signature:
+    "Tensor (Tensor a, Tensor b, Int64List dims_a, Int64List dims_b) => TensorDot"
+  bind_python: True
+
 - name: "l1_loss"
   signature: "Tensor(Tensor input, Tensor target, String reduction) => L1Loss"
   bind_python: True

oneflow/core/functional/impl/nn_functor.cpp

@@ -264,6 +264,93 @@ class BatchMatMulFunctor {
   std::shared_ptr<OpExpr> batch_matmul_op_;
 };
 
+class TensorDotFunctor {
+ public:
+  Maybe<Tensor> operator()(const std::shared_ptr<Tensor>& a, const std::shared_ptr<Tensor>& b,
+                           const std::vector<int64_t>& dims_a,
+                           const std::vector<int64_t>& dims_b) const {
+    CHECK_EQ_OR_RETURN(dims_a.size(), dims_b.size())
+        << "dims1 and dims2 must have same size, got " << dims_a.size() << " and " << dims_b.size();
+
+    if (dims_a.empty() && dims_b.empty()) {
+      std::vector<int64_t> shape_sum;
+      shape_sum.reserve(a->shape()->NumAxes() + b->shape()->NumAxes());
+      for (int64_t i = 0; i < a->shape()->NumAxes(); i++) {
+        shape_sum.emplace_back(a->shape()->At(i));
+      }
+      for (int64_t i = 0; i < b->shape()->NumAxes(); i++) {
+        shape_sum.emplace_back(b->shape()->At(i));
+      }
+      auto reshape_a = JUST(Reshape(a, Shape(DimVector{-1, 1})));
+      auto reshape_b = JUST(Reshape(b, Shape(DimVector{1, -1})));
+      return JUST(Reshape(JUST(functional::MatMul(reshape_a, reshape_b, false, false, 1.0)),
+                          Shape(DimVector(shape_sum.begin(), shape_sum.end()))));
+    }
+    std::vector<bool> if_dot_dims_a(a->shape()->NumAxes(), false);
+    std::vector<bool> if_dot_dims_b(b->shape()->NumAxes(), false);
+    for (auto i : dims_a) if_dot_dims_a[i] = true;
+    for (auto i : dims_b) if_dot_dims_b[i] = true;
+
+    std::vector<int32_t> broadcast_dims_a, broadcast_dims_b;
+    for (int64_t i = 0; i < dims_a.size(); i++) {
+      int64_t size_a = a->shape()->At(dims_a[i]);
+      int64_t size_b = b->shape()->At(dims_b[i]);
+      if (size_a == 1) {
+        broadcast_dims_b.emplace_back(dims_b[i]);
+      } else if (size_b == 1) {
+        broadcast_dims_a.emplace_back(dims_a[i]);
+      } else {
+        CHECK_EQ_OR_RETURN(size_a, size_b) << "The corresponding dim must be equal, got " << size_a
+                                           << " in tensor a and " << size_b << " in tensor b";
+      }
+    }
+    auto reduced_a = a;
+    auto reduced_b = b;
+    if (!broadcast_dims_a.empty())
+      reduced_a = JUST(functional::ReduceSum(a, broadcast_dims_a, true));
+    if (!broadcast_dims_b.empty())
+      reduced_b = JUST(functional::ReduceSum(b, broadcast_dims_b, true));
+
+    int64_t rsize_a = 1, rsize_b = 1;
+    std::vector<int64_t> rshape_a, rshape_b;
+    std::vector<int32_t> pa, pb;
+    pa.reserve(a->shape()->NumAxes());
+    pb.reserve(b->shape()->NumAxes());
+    std::vector<int64_t> non_dot_dims;
+    for (int64_t i = 0; i < a->shape()->NumAxes(); i++) {
+      if (!if_dot_dims_a[i]) {
+        non_dot_dims.emplace_back(a->shape()->At(i));
+        pa.emplace_back(i);
+        rsize_a *= reduced_a->shape()->At(i);
+        rshape_a.emplace_back(reduced_a->shape()->At(i));
+      }
+    }
+
+    for (auto i : dims_a) pa.emplace_back(i);
+    for (auto i : dims_b) pb.emplace_back(i);
+
+    for (int64_t i = 0; i < b->shape()->NumAxes(); i++) {
+      if (!if_dot_dims_b[i]) {
+        non_dot_dims.emplace_back(b->shape()->At(i));
+        pb.emplace_back(i);
+        rsize_b *= reduced_b->shape()->At(i);
+        rshape_b.emplace_back(reduced_b->shape()->At(i));
+      }
+    }
+    rshape_a.insert(rshape_a.end(), rshape_b.begin(), rshape_b.end());
+
+    int64_t dshape = 1;
+    for (auto i : dims_a) dshape *= reduced_a->shape()->At(i);
+    auto permute_a =
+        JUST(Reshape(JUST(Permute(reduced_a, pa)), Shape(DimVector({rsize_a, dshape}))));
+    auto permute_b =
+        JUST(Reshape(JUST(Permute(reduced_b, pb)), Shape(DimVector({dshape, rsize_b}))));
+
+    return Reshape(JUST(functional::MatMul(permute_a, permute_b, false, false, 1.0)),
+                   Shape(DimVector({rshape_a.begin(), rshape_a.end()})));
+  }
+};
+
 class FusedMLPFunctor {
  public:
   FusedMLPFunctor() {
@@ -2721,6 +2808,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::DeConv3dFunctor>("Deconv3d");
   m.add_functor<impl::MatMulFunctor>("MatMul");
   m.add_functor<impl::BatchMatMulFunctor>("BatchMatMul");
+  m.add_functor<impl::TensorDotFunctor>("TensorDot");
   m.add_functor<impl::FusedMLPFunctor>("FusedMLP");
   m.add_functor<impl::LayerNormFunctor>("LayerNorm");
   m.add_functor<impl::LayerNormAffineFunctor>("LayerNormAffine");

oneflow/ir/include/OneFlow/OneFlowUserOps.td

@@ -211,7 +211,7 @@ class OneFlow_NormalizationAddReluBaseOp : OneFlow_BaseOp<"normalization_add_rel
 #endif // GET_ONEFLOW_BASE_OP_DEFINITIONS
 
 // Group: BINARY
-// bias_add, cast_like, celu_grad, diag_grad, diagonal_grad, dot, dropout_grad, elementwise_maximum, elementwise_minimum, elu_grad, floordiv, gelu_grad, grid_sample, hardsigmoid_grad, hardshrink_grad, hardswish_grad, l1_l2_regularize_gradient, leaky_relu_grad, masked_fill, mish_grad, multiply, narrow_grad, pow, prelu, relu_grad, selu_grad, sigmoid_grad, silu_grad, softshrink_grad, threshold_grad, tf_prelu, unfold_tensor_grad, xdivy, xlogy
+// bias_add, cast_like, celu_grad, diag_grad, diagonal_grad, dot, dropout_grad, elementwise_maximum, elementwise_minimum, elu_grad, floordiv, gelu_grad, grid_sample, hardsigmoid_grad, hardshrink_grad, hardswish_grad, l1_l2_regularize_gradient, leaky_relu_grad, masked_fill, mish_grad, multiply, narrow_grad, pow, prelu, relu_grad, selu_grad, sigmoid_grad, silu_grad, softshrink_grad, tensor_dot, threshold_grad, tf_prelu, unfold_tensor_grad, xdivy, xlogy
 // Total: 34
 
 #ifdef GET_ONEFLOW_BINARY_OP_DEFINITIONS
@@ -659,6 +659,24 @@ def OneFlow_SiluGradOp : OneFlow_BaseOp<"silu_grad", [NoSideEffect, DeclareOpInt
   let has_data_type_infer_fn = 1;
 }
 
+def OneFlow_TensorDotOp : OneFlow_BaseOp<"tensordot", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
+  let input = (ins
+    OneFlow_Tensor:$x,
+    OneFlow_Tensor:$y
+  );
+  let output = (outs
+    OneFlow_Tensor:$z
+  );
+  let attrs = (ins
+    DefaultValuedAttr<SI32Attr, "0">:$axis_a,
+    DefaultValuedAttr<SI32Attr, "0">:$axis_b
+  );
+  let has_logical_tensor_desc_infer_fn = 1;
+  let has_physical_tensor_desc_infer_fn = 1;
+  let has_get_sbp_fn = 1;
+  let has_data_type_infer_fn = 1;
+}
+
 def OneFlow_ThresholdGradOp : OneFlow_BaseOp<"threshold_grad", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
   let input = (ins
     OneFlow_Tensor:$x,

python/oneflow/__init__.py

@@ -364,6 +364,7 @@ def atexit_hook(hook):
 from oneflow.nn.modules.global_cast import to_local_op as to_local
 from oneflow.nn.modules.where import where_op as where
 from oneflow.nn.modules.scatter import *
+from oneflow.nn.functional import tensordot
 from oneflow.ops.stateful_ops import StatefulOp as stateful_op
 from oneflow.ops.initializer_util import constant_initializer
 from oneflow.ops.initializer_util import glorot_normal_initializer

python/oneflow/framework/docstr/__init__.py

@@ -68,3 +68,4 @@
 from .argsort import *
 from .module import *
 from .util_ops import *
+from .tensordot import *

python/oneflow/framework/docstr/tensordot.py

@@ -0,0 +1,54 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import oneflow
+from oneflow.framework.docstr.utils import add_docstr
+
+add_docstr(
+    oneflow.tensordot,
+    r"""
+    tensordot(a, b, dims=Union[int, Tensor, Tuple]) -> Tensor
+
+    Compute tensor dot along given dimensions.
+
+    Given two tensors a and b, and dims which has two list containing dim indices, `tensordot` traverses the two 
+    lists and calculate the tensor dot along every dim pair.
+
+    Args:
+        a: The input tensor to compute tensordot
+        b: The input tensor to compute tensordot
+        dims: int or array-like
+
+    Returns:
+        Oneflow.Tensor: The result tensor
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> a = flow.randn(3, 4, 5)
+        >>> b = flow.randn(4, 5, 6)
+        >>> flow.tensordot(a, b, dims=2).shape
+        oneflow.Size([3, 6])
+        >>> b = flow.randn(5, 6, 7)
+        >>> flow.tensordot(a, b, dims=1).shape
+        oneflow.Size([12, 42])
+        >>> b = flow.randn(3, 4, 7)
+        >>> flow.tensordot(a, b, dims=[[0, 1], [0, 1]]).shape
+        oneflow.Size([5, 7])
+
+    """,
+)

python/oneflow/nn/functional/__init__.py

@@ -26,6 +26,7 @@
 from .functional_maxpool import max_pool1d
 from .functional_maxpool import max_pool2d
 from .functional_maxpool import max_pool3d
+from .functional_tensordot import tensordot
 from oneflow._C import adaptive_avg_pool1d
 from oneflow._C import adaptive_avg_pool2d
 from oneflow._C import adaptive_avg_pool3d

python/oneflow/nn/functional/functional_tensordot.py

@@ -0,0 +1,40 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import oneflow
+from typing import Union, List, Tuple
+
+def tensordot(a, b, dims: Union[int, List[List[int]]] = 2):
+    if not isinstance(dims, (oneflow._oneflow_internal.Tensor, int, list, tuple)):
+        raise TypeError(
+            f"oneflow.tensordot expects dims to be one of oneflow.Tensor, int, List[List[int]] or Tuple[List[int], List[int]], but got {type(dims)}"
+        )
+
+    if isinstance(dims, int):
+        assert dims >= 0 and dims <= min(a.dim(), b.dim())
+        dim_a = list(range(a.dim() - dims, a.dim()))
+        dim_b = list(range(dims))
+
+    elif isinstance(dims, (list, tuple, oneflow._oneflow_internal.Tensor)):
+        assert len(dims) == 2
+        dim_a = list(dims[0])
+        dim_b = list(dims[1])
+        assert (
+            isinstance(dim_a[0], int)
+            and isinstance(dim_b[0], int)
+            and len(dim_a) == len(dim_b)
+        )
+
+    return oneflow._C.tensordot(a, b, dim_a, dim_b)

python/oneflow/test/modules/test_tensordot.py

@@ -0,0 +1,64 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import unittest
+import numpy as np
+import oneflow as flow
+import oneflow.unittest
+from oneflow.test_utils.automated_test_util import *
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestTensordot(flow.unittest.TestCase):
+    @autotest(check_graph=True)
+    def test_tensordot_intdim(test_case):
+        device = random_device()
+        dims = random()
+        dims_list = [random().to(int).value() for i in range(dims.to(int).value() + 3)]
+        x = random_tensor(
+            ndim=3,
+            dim0=dims_list[0],
+            dim1=dims_list[1],
+            dim2=dims_list[2],
+        ).to(device)
+        y = random_tensor(
+            ndim=3,
+            dim0=dims_list[0 + dims.to(int).value()],
+            dim1=dims_list[1 + dims.to(int).value()],
+            dim2=dims_list[2 + dims.to(int).value()],
+        ).to(device)
+
+        z = torch.tensordot(x, y, dims=3 - dims.to(int).value())
+        return z
+
+    @autotest(check_graph=True, n=1)
+    def test_tensordot_tensordim(test_case):
+        device = random_device()
+        x = random_tensor(4, 1,3,2,5).to(device)
+        y = random_tensor(4, 4,2,3,5).to(device)
+        z = torch.tensordot(x, y, dims=[[1,2,0],[2,1,0]])
+        return z
+
+    @autotest(check_graph=True)
+    def test_tensordot_broadcast(test_case):
+        device = random_device()
+        x = random_tensor(4, 1,1,1,1).to(device)
+        y = random_tensor(4, 2,3,4,5).to(device)
+        z = torch.tensordot(x, y, dims=random(high=5).to(int).value())
+        return z
+
+
+if __name__ == "__main__":
+    unittest.main()