Migrate tile python layer to functor (#7305)

* tile implement

* migrate repeat

* of format

* align document with pytorch

* change input args to *size

* change with review comments

* migrate tile python layer to functor

* fix tile document

* fix tile document

* add document's function signature

* auto format by CI

* add repeat document signature

* fix document

* fix document

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>

oneflow/core/functional/functional_api.yaml

@@ -748,6 +748,14 @@
   signature: "Tensor (Tensor dy, Int32List out, Int32List expand) => ExpandGrad"
   bind_python: False
 
+- name: "repeat"
+  signature: "Tensor (Tensor input, Shape repeat_shape) => Repeat"
+  bind_python: True
+
+- name: "tile"
+  signature: "Tensor (Tensor input, Shape dims) => Tile"
+  bind_python: True
+
 - name: "roll"
   signature: "Tensor (Tensor x, Int32List[1] shifts, Int32List[1] dims=None) => Roll"
   bind_python: True

oneflow/core/functional/impl/array_functor.cpp

@@ -2561,6 +2561,75 @@ class GenTensorBufferFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class RepeatFunctor {
+ public:
+  RepeatFunctor() {}
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& input,
+                           const Shape& repeat_shape) const {
+    Shape input_shape = *(input->shape());
+    std::vector<int32_t> input_reshape_vec;
+    std::vector<int32_t> expand_shape_vec;
+    std::vector<int32_t> output_reshape_vec;
+
+    int32_t numaxes_diff = repeat_shape.NumAxes() - input_shape.NumAxes();
+    CHECK_GE_OR_RETURN(numaxes_diff, 0)
+        << "RuntimeError: Number of dimensions of repeat dims can not be "
+           "smaller than number of dimensions of tensor";
+
+    for (int32_t i = repeat_shape.NumAxes() - 1; i >= 0; i--) {
+      if (i >= numaxes_diff) {
+        int32_t input_shape_val = input_shape.At(i - numaxes_diff);
+        int32_t repeat_shape_val = repeat_shape.At(i);
+        if (repeat_shape_val > 1) {
+          if (input_shape_val > 1) {
+            input_reshape_vec.insert(input_reshape_vec.begin(), input_shape_val);
+            input_reshape_vec.insert(input_reshape_vec.begin(), 1);
+            expand_shape_vec.insert(expand_shape_vec.begin(), input_shape_val);
+            expand_shape_vec.insert(expand_shape_vec.begin(), repeat_shape_val);
+            output_reshape_vec.insert(output_reshape_vec.begin(),
+                                      repeat_shape_val * input_shape_val);
+          } else {
+            input_reshape_vec.insert(input_reshape_vec.begin(), input_shape_val);
+            expand_shape_vec.insert(expand_shape_vec.begin(), repeat_shape_val);
+            output_reshape_vec.insert(output_reshape_vec.begin(), repeat_shape_val);
+          }
+        } else {
+          input_reshape_vec.insert(input_reshape_vec.begin(), input_shape_val);
+          expand_shape_vec.insert(expand_shape_vec.begin(), input_shape_val);
+          output_reshape_vec.insert(output_reshape_vec.begin(), input_shape_val);
+        }
+      } else {
+        expand_shape_vec.insert(expand_shape_vec.begin(), repeat_shape.At(i));
+        output_reshape_vec.insert(output_reshape_vec.begin(), repeat_shape.At(i));
+      }
+    }
+    Shape input_reshape(DimVector(input_reshape_vec.begin(), input_reshape_vec.end()));
+    Shape expand_shape(DimVector(expand_shape_vec.begin(), expand_shape_vec.end()));
+    Shape output_reshape(DimVector(output_reshape_vec.begin(), output_reshape_vec.end()));
+    std::shared_ptr<one::Tensor> reshaped_tensor = JUST(Reshape(input, input_reshape));
+    std::shared_ptr<one::Tensor> expanded_tensor = JUST(Expand(reshaped_tensor, expand_shape));
+    std::shared_ptr<one::Tensor> result = JUST(Reshape(expanded_tensor, output_reshape));
+    return result;
+  }
+};
+
+class TileFunctor {
+ public:
+  TileFunctor() {}
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& input, const Shape& dims) const {
+    std::vector<int32_t> new_dims_vec;
+    int32_t numaxes_diff = input->shape()->NumAxes() - dims.NumAxes();
+    for (int32_t i = dims.NumAxes() - 1; i >= 0; i--) {
+      CHECK_GE_OR_RETURN(dims.At(i), 0)
+          << "RuntimeError: Tring to create tensor with negative dimension " << dims.At(i);
+      new_dims_vec.insert(new_dims_vec.begin(), dims.At(i));
+    }
+    for (int32_t i = 0; i < numaxes_diff; i++) { new_dims_vec.insert(new_dims_vec.begin(), 1); }
+    Shape new_dims(DimVector(new_dims_vec.begin(), new_dims_vec.end()));
+    return JUST(Repeat(input, new_dims));
+  }
+};
+
 class TransposeAllDimPropertyFunctor {
  public:
   TransposeAllDimPropertyFunctor() {}
@@ -2690,6 +2759,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::TensorToTensorBufferFunctor>("TensorToTensorBuffer");
   m.add_functor<impl::TensorBufferToTensorFunctor>("TensorBufferToTensor");
   m.add_functor<impl::GenTensorBufferFunctor>("GenTensorBuffer");
+  m.add_functor<impl::RepeatFunctor>("Repeat");
+  m.add_functor<impl::TileFunctor>("Tile");
   m.add_functor<impl::TransposeAllDimPropertyFunctor>("TransposeAllDimProperty");
   m.add_functor<impl::TransposeAllDimFunctionFunctor>("TransposeAllDimFunction");
 };

python/oneflow/__init__.py

@@ -94,6 +94,8 @@ def is_deprecated(func_or_class):
 from oneflow._C import logical_not
 from oneflow._C import gelu
 from oneflow._C import mish
+from oneflow._C import repeat
+from oneflow._C import tile
 from oneflow._C import sigmoid
 from oneflow._C import tanh
 from oneflow._C import silu
@@ -342,7 +344,6 @@ def atexit_hook(hook):
 from oneflow.nn.modules.reduce_ops import prod_op as prod
 from oneflow.nn.modules.reduce_ops import all_op as all
 from oneflow.nn.modules.reduce_ops import any_op as any
-from oneflow.nn.modules.repeat import repeat_op as repeat
 from oneflow.nn.modules.reshape import reshape_op as reshape
 from oneflow.nn.modules.reshape import view_op as view
 from oneflow.nn.modules.slice import slice_op as slice
@@ -356,7 +357,6 @@ def atexit_hook(hook):
 )
 from oneflow.nn.modules.as_tensor import as_tensor
 from oneflow.nn.modules.tensor_buffer import tensor_to_tensor_buffer
-from oneflow.nn.modules.tile import tile_op as tile
 from oneflow.nn.modules.consistent_cast import to_consistent_op as to_consistent
 from oneflow.nn.modules.consistent_cast import to_local_op as to_local
 from oneflow.nn.modules.where import where_op as where

python/oneflow/framework/docstr/__init__.py

@@ -45,4 +45,6 @@
 from .clamp import *
 from .erfinv import *
 from .swapaxes import *
+from .repeat import *
+from .tile import *
 from .tensor_t import *

python/oneflow/framework/docstr/repeat.py

@@ -0,0 +1,48 @@
+"""
+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.repeat,
+    """
+    repeat(input, sizes) -> Tensor
+
+    This operator repeat the input tensor to a larger size along the specified dimensions.
+
+    Args:
+        input (oneflow.Tensor): The input Tensor.
+        sizes (flow.Shape or List): The number of times to repeat this tensor along each dimension.
+
+    Returns:
+        oneflow.Tensor: The result Tensor.
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> import numpy as np
+        >>> np_arr = np.random.randn(5, 3, 6, 9).astype(np.float32)
+        >>> input = flow.Tensor(np_arr)
+        >>> out = input.repeat(1, 1, 2, 2)
+        >>> out.shape
+        oneflow.Size([5, 3, 12, 18])
+        >>> out = input.repeat(2, 1, 1, 2, 2)
+        >>> out.shape
+        oneflow.Size([2, 5, 3, 12, 18])
+    """,
+)

python/oneflow/framework/docstr/tensor.py

@@ -841,6 +841,15 @@
     """,
 )
 
+add_docstr(
+    oneflow.Tensor.repeat,
+    """
+    Tensor.repeat(*size) -> Tensor
+
+    See :func:`oneflow.repeat`
+    """,
+)
+
 add_docstr(
     oneflow.Tensor.t,
     """
@@ -850,6 +859,15 @@
     """,
 )
 
+add_docstr(
+    oneflow.Tensor.tile,
+    """
+    Tensor.tile(*dims) -> Tensor
+
+    See :func:`oneflow.tile`
+    """,
+)
+
 add_docstr(
     oneflow.Tensor.T,
     """

python/oneflow/nn/modules/tile.py → python/oneflow/framework/docstr/tile.py

@@ -13,36 +13,36 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from typing import Union
-
-import oneflow as flow
-from oneflow.framework.tensor import register_tensor_op
+import oneflow
+from oneflow.framework.docstr.utils import add_docstr
 
+add_docstr(
+    oneflow.tile,
+    """
+    tile(input, dims) -> Tensor
 
-@register_tensor_op("tile")
-def tile_op(input, reps):
-    """The interface is consistent with PyTorch.
+    The interface is consistent with PyTorch.
     The documentation is referenced from:
     https://pytorch.org/docs/stable/generated/torch.tile.html
 
-    Constructs a tensor by repeating the elements of ``input``.  The ``reps`` argument specifies the number
+    Constructs a tensor by repeating the elements of ``input``.  The ``dims`` argument specifies the number
     of repetitions in each dimension.
 
-    If ``reps`` specifies fewer dimensions than ``input`` has, then ones are prepended to ``reps`` until
-    all dimensions are specified.  For example, if ``input`` has shape (8, 6, 4, 2) and ``reps`` is (2, 2),
-    then ``reps`` is treated as (1, 1, 2, 2).
+    If ``dims`` specifies fewer dimensions than ``input`` has, then ones are prepended to ``dims`` until
+    all dimensions are specified.  For example, if ``input`` has shape (8, 6, 4, 2) and ``dims`` is (2, 2),
+    then ``dims`` is treated as (1, 1, 2, 2).
 
-    Analogously, if ``input`` has fewer dimensions than ``reps`` specifies, then ``input`` is treated as
-    if it were unsqueezed at dimension zero until it has as many dimensions as ``reps`` specifies.
-    For example, if ``input`` has shape (4, 2) and ``reps`` is (3, 3, 2, 2), then ``input`` is treated as
+    Analogously, if ``input`` has fewer dimensions than ``dims`` specifies, then ``input`` is treated as
+    if it were unsqueezed at dimension zero until it has as many dimensions as ``dims`` specifies.
+    For example, if ``input`` has shape (4, 2) and ``dims`` is (3, 3, 2, 2), then ``input`` is treated as
     if it had the shape (1, 1, 4, 2).
 
     .. note::
         This function is similar to NumPy’s tile function.
 
     Args:
         input (oneflow.Tensor): the tensor whose elements to repeat.
-        reps (tuple): the number of repetitions per dimension.
+        dims (tuple): the number of repetitions per dimension.
 
     For example:
 
@@ -51,32 +51,15 @@ def tile_op(input, reps):
         >>> import oneflow as flow
         >>> import numpy as np
         
-        >>> x = np.array([1, 2]).astype(np.int32)
-        >>> input = flow.tensor(x, dtype=flow.int32)
-        >>> out = input.tile(reps=(2,))
-        >>> out
-        tensor([1, 2, 1, 2], dtype=oneflow.int32)
-
+        >>> np_arr = np.random.randn(5, 3, 6, 9).astype(np.float32)
+        >>> input = flow.Tensor(np_arr)
+        >>> out = input.tile(2,1,2,1)
+        >>> out.shape
+        oneflow.Size([10, 3, 12, 9])
         >>> x = np.random.randn(5, 2, 1)
         >>> input = flow.Tensor(x)
-        >>> out = input.tile(reps=(3, 4))
-        >>> out.size()
+        >>> out = input.tile(3,4)
+        >>> out.shape
         oneflow.Size([5, 6, 4])
-
-    """
-
-    for s in reps:
-        assert s > 0
-    input_shape = input.shape
-    diff = len(input_shape) - len(reps)
-    if diff > 0:
-        shape = [1 for _ in range(diff)]
-        shape.extend([i for i in reps])
-        reps = tuple(shape)
-    return input.repeat(reps)
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod(raise_on_error=True)
+    """,
+)

python/oneflow/framework/tensor.py

@@ -753,6 +753,26 @@ def _gather(self, dim, index):
     return flow._C.dim_gather(self, dim, index, False)
 
 
+def _repeat(self, *sizes):
+    if len(sizes) == 1:
+        new_sizes = sizes[0]
+        if isinstance(new_sizes, int):
+            new_sizes = (new_sizes,)
+    else:
+        new_sizes = sizes
+    return flow._C.repeat(self, new_sizes)
+
+
+def _tile(self, *dims):
+    if len(dims) == 1:
+        new_dims = dims[0]
+        if isinstance(new_dims, int):
+            new_dims = (new_dims,)
+    else:
+        new_dims = dims
+    return flow._C.tile(self, new_dims)
+
+
 def _T(self):
     return flow._C.T(self)
 
@@ -921,6 +941,8 @@ def RegisterMethods():
     Tensor.roll = _roll
     Tensor.bmm = _bmm
     Tensor.chunk = _chunk
+    Tensor.repeat = _repeat
+    Tensor.tile = _tile
     Tensor.split = _split
     Tensor.squeeze = _squeeze
     Tensor.swapaxes = _swapaxes