Add mv op

docs/source/oneflow.rst

@@ -92,6 +92,7 @@ oneflow
             masked_fill, 
             masked_select, 
             matmul, 
+            mv, 
             narrow, 
             max, 
             mean,

docs/source/tensor.rst

@@ -104,6 +104,7 @@ OneFlow Tensor Class
             masked_fill, 
             masked_select, 
             matmul, 
+            mv, 
             max, 
             mean, 
             min, 

oneflow/core/functional/functional_api.yaml

@@ -968,6 +968,11 @@
     Double alpha=1.0) => MatMul"
   bind_python: True
 
+- name: "mv"
+  signature:
+    "Tensor (Tensor input, Tensor vec) => Mv"
+  bind_python: True
+
 - name: "fused_mlp"
   signature:
     "Tensor (Tensor x, TensorTuple weights, TensorTuple biases, Bool skip_final_activation) => FusedMLP"

oneflow/core/functional/impl/nn_functor.cpp

@@ -3340,6 +3340,29 @@ class RocAucScoreFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class MvFunctor {
+ public:
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& input,
+                           const std::shared_ptr<one::Tensor>& vec) const {
+    const auto& input_shape = input->shape();
+    const auto& vec_shape = vec->shape();
+    CHECK_OR_RETURN(input_shape->NumAxes() == 2 && vec_shape->NumAxes() == 1)
+        << Error::RuntimeError() << "vector + matrix @ vector expected, got "
+        << "1, " << input_shape->NumAxes() << ", " << vec_shape->NumAxes();
+    CHECK_EQ_OR_RETURN(input_shape->at(1), vec_shape->at(0))
+        << Error::RuntimeError() << "size mismatch, got " << std::to_string(input_shape->at(0))
+        << ", " << std::to_string(input_shape->at(0)) << "x" << std::to_string(input_shape->at(1))
+        << ", " << std::to_string(vec_shape->at(0));
+    // TODO(zhongshsh): speedup
+    const std::shared_ptr<Tensor> reshape_vec =
+        JUST(Reshape(vec, Shape(DimVector{vec_shape->at(0), 1})));
+    std::shared_ptr<Tensor> out = JUST(MatMul(input, reshape_vec, false, false, 1.0));
+    std::shared_ptr<Tensor> reshape_out = JUST(Squeeze(
+        JUST(Reshape(out, Shape(DimVector{1, input_shape->at(0)}))), std::vector<int32_t>({0})));
+    return reshape_out;
+  }
+};
+
 }  // namespace impl
 
 ONEFLOW_FUNCTION_LIBRARY(m) {
@@ -3353,6 +3376,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::EmbeddingReNormFunctor>("EmbeddingReNorm");
   m.add_functor<impl::EmbeddingFunctor>("Embedding");
   m.add_functor<impl::MatMulFunctor>("MatMul");
+  m.add_functor<impl::MvFunctor>("Mv");
   m.add_functor<impl::BatchMatMulFunctor>("BatchMatMul");
   m.add_functor<impl::TensorDotFunctor>("TensorDot");
   m.add_functor<impl::TensorDotIntDimsFunctor>("TensorDotIntDims");

python/oneflow/__init__.py

@@ -152,6 +152,7 @@ def is_deprecated(func_or_class):
 from oneflow._C import sqrt
 from oneflow._C import square
 from oneflow._C import matmul
+from oneflow._C import mv
 from oneflow._C import bernoulli
 from oneflow._C import round
 from oneflow._C import softplus

python/oneflow/framework/docstr/math_ops.py

@@ -1274,6 +1274,39 @@
     """,
 )
 
+add_docstr(
+    oneflow.mv,
+    r"""
+    mv(input, vec) -> Tensor
+
+    The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.mv.html.
+
+    Performs a matrix-vector product of the matrix :attr:`input` and the vector :attr:`vec`.
+
+    If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`vec` is a
+    1-D tensor of size `m`, :attr:`out` will be a 1-D tensor of size `n`.
+
+    .. note:: This function does not :ref:`broadcast <broadcasting-semantics>`.
+
+    Args:
+        input (oneflow.Tensor): matrix to be matrix multiplied
+        vec (oneflow.Tensor): vector to be matrix multiplied
+    Returns:
+        oneflow.Tensor: the output Tensor
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> mat = flow.randn(2, 3)
+        >>> vec = flow.randn(3)
+        >>> out = flow.mv(mat, vec)
+        >>> out.shape
+        oneflow.Size([2])
+    """,
+)
+
 add_docstr(
     oneflow.round,
     r"""This operator rounds the value of Blob to the nearest integer.

python/oneflow/framework/docstr/tensor.py

@@ -603,6 +603,13 @@
     """,
 )
 
+add_docstr(
+    oneflow.Tensor.mv,
+    """
+    See :func:`oneflow.mv`
+    """,
+)
+
 add_docstr(
     oneflow.Tensor.narrow,
     """

python/oneflow/framework/tensor.py

@@ -601,6 +601,10 @@ def _matmul(self, other):
     return flow.matmul(self, other)
 
 
+def _mv(self, vec):
+    return flow._C.mv(self, vec)
+
+
 def _round(self):
     return flow.round(self)
 
@@ -1152,6 +1156,7 @@ def RegisterMethods():
     Tensor.new_tensor = _new_tensor
     Tensor.cumsum = _cumsum
     Tensor.cumprod = _cumprod
+    Tensor.mv = _mv
 
 
 def register_tensor_op(op_name):

python/oneflow/test/exceptions/test_mv.py

@@ -0,0 +1,50 @@
+"""
+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 oneflow as flow
+import oneflow.unittest
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestMv(flow.unittest.TestCase):
+    def test_mv_not_matrix(test_case):
+        with test_case.assertRaises(Exception) as exp:
+            mat = flow.randn(2, 3, 3)
+            vec = flow.randn(3)
+            out = flow.mv(mat, vec)
+        test_case.assertTrue(
+            "vector + matrix @ vector expected, got 1, 3, 1" in str(exp.exception)
+        )
+
+    def test_mv_not_vector(test_case):
+        with test_case.assertRaises(Exception) as exp:
+            mat = flow.randn(2, 3)
+            vec = flow.randn(3, 1)
+            out = flow.mv(mat, vec)
+        test_case.assertTrue(
+            "vector + matrix @ vector expected, got 1, 2, 2" in str(exp.exception)
+        )
+
+    def test_mv_size_mismatch(test_case):
+        with test_case.assertRaises(Exception) as exp:
+            mat = flow.randn(2, 3)
+            vec = flow.randn(4)
+            out = flow.mv(mat, vec)
+        test_case.assertTrue("size mismatch" in str(exp.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()

python/oneflow/test/modules/test_consistent_mv.py

@@ -0,0 +1,39 @@
+"""
+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 oneflow as flow
+import oneflow.unittest
+from oneflow.test_utils.automated_test_util import *
+
+
+@autotest(n=1, check_graph=False)
+def _test_mv(test_case, placement, sbp):
+    dim = random(1, 6)
+    mat = random_tensor(2, dim1=dim).to_global(placement=placement, sbp=sbp)
+    vec = random_tensor(1, dim0=dim).to_global(placement=placement, sbp=sbp)
+    return torch.mv(mat, vec)
+
+
+class TestMvModule(flow.unittest.TestCase):
+    @globaltest
+    def test_mv(test_case):
+        for placement in all_placement():
+            for sbp in all_sbp(placement):
+                _test_mv(test_case, placement, sbp)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/oneflow/test/modules/test_matmul.py

@@ -49,6 +49,19 @@ def test_flow_tensor_broadcast_matmul_with_random_data(test_case):
         y = random_tensor(ndim=2, dim0=k).to(device)
         return x.matmul(y)
 
+    @autotest(check_graph=True)
+    def test_flow_mv_with_random_data(test_case):
+        device = random_device()
+        k = random(1, 6)
+        x = random_tensor(ndim=2, dim1=k).to(device)
+        y = random_tensor(ndim=1, dim0=k).to(device)
+        z = torch.mv(x, y)
+        return z
+
+    @profile(torch.mv)
+    def profile_mv(test_case):
+        torch.mv(torch.ones(32, 64), torch.ones(64))
+
 
 if __name__ == "__main__":
     unittest.main()

python/oneflow/test/tensor/test_tensor_part_1.py

@@ -416,6 +416,16 @@ def test_matmul_with_random_data(test_case):
         b = random_tensor(ndim=2, dim0=dim1, dim1=dim2)
         return a @ b
 
+    @flow.unittest.skip_unless_1n1d()
+    @autotest()
+    def test_mm_with_random_data(test_case):
+        device = random_device()
+        dim0 = random(low=2, high=10).to(int)
+        dim1 = random(low=3, high=20).to(int)
+        a = random_tensor(ndim=2, dim0=dim0, dim1=dim1).to(device)
+        b = random_tensor(ndim=1, dim0=dim1).to(device)
+        return a.mv(b)
+
     @flow.unittest.skip_unless_1n1d()
     def test_tensor_to_list(test_case):
         list_data = [[1.0, 3.0], [5.0, 6.0]]