!20011 [assistant][ops] Add nn operator Roll

Merge pull request !20011 from 孟权令/Roll

mindspore/core/base/core_ops.h

@@ -93,6 +93,7 @@ constexpr auto kDropoutDoMask = "DropoutDoMask";
 constexpr auto kDropout = "Dropout";
 constexpr auto kDropoutGrad = "DropoutGrad";
 constexpr auto kConv2DTranspose = "Conv2DTranspose";
+constexpr auto kRoll = "Roll";
 
 // Here list all primitives used in backend or some special primitives used by core.
 // GetNext
@@ -283,6 +284,7 @@ inline const PrimitivePtr kPrimCTCLossV2Grad = std::make_shared<Primitive>("CTCL
 inline const PrimitivePtr kPrimCTCLoss = std::make_shared<Primitive>(kCTCLoss);
 inline const PrimitivePtr kPrimFullConnection = std::make_shared<Primitive>("FullConnection");
 inline const PrimitivePtr kPrimConv2DTranspose = std::make_shared<Primitive>(kConv2DTranspose);
+inline const PrimitivePtr kPrimRoll = std::make_shared<Primitive>(kRoll);
 inline const PrimitivePtr kPrimGroupConv2DGradInput = std::make_shared<Primitive>("GroupConv2DGradInput");
 inline const PrimitivePtr kPrimBatchNorm = std::make_shared<Primitive>("BatchNorm");
 inline const PrimitivePtr kPrimBatchNormGrad = std::make_shared<Primitive>("BatchNormGrad");

mindspore/core/ops/roll.cc

@@ -0,0 +1,56 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+
+#include "ops/roll.h"
+#include <set>
+
+#include "ops/op_utils.h"
+#include "utils/check_convert_utils.h"
+#include "utils/tensor_construct_utils.h"
+#include "abstract/primitive_infer_map.h"
+
+namespace mindspore {
+namespace ops {
+namespace {
+abstract::ShapePtr InferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  auto prim_name = primitive->name();
+  CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kEqual, 1, prim_name);
+  auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
+  auto axis = GetValue<int64_t>(primitive->GetAttr(kAxis));
+  auto x_rank = SizeToLong(x_shape.size());
+  CheckAndConvertUtils::CheckInRange("axis value", axis, kIncludeLeft, {-x_rank, x_rank}, prim_name);
+  return std::make_shared<abstract::Shape>(x_shape);
+}
+
+TypePtr InferType(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &input_args) {
+  for (const auto &item : input_args) {
+    MS_EXCEPTION_IF_NULL(item);
+  }
+  const std::set<TypePtr> valid_types = {kFloat32, kFloat16, kInt32, kUInt32, kInt8, kUInt8};
+  auto infer_type = input_args[0]->BuildType();
+  return CheckAndConvertUtils::CheckTensorTypeValid("x type", infer_type, valid_types, prim->name());
+}
+}  // namespace
+
+AbstractBasePtr RollInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                          const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  return abstract::MakeAbstract(InferShape(primitive, input_args), InferType(primitive, input_args));
+}
+REGISTER_PRIMITIVE_EVAL_IMPL(Roll, prim::kPrimRoll, RollInfer, nullptr, true);
+}  // namespace ops
+}  // namespace mindspore

mindspore/core/ops/roll.h

@@ -0,0 +1,41 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+
+#ifndef MINDSPORE_CORE_OPS_ROLL_H_
+#define MINDSPORE_CORE_OPS_ROLL_H_
+#include <memory>
+#include <vector>
+
+#include "ops/primitive_c.h"
+#include "abstract/abstract_value.h"
+#include "utils/check_convert_utils.h"
+
+namespace mindspore {
+namespace ops {
+constexpr auto kNameRoll = "Roll";
+class Roll : public PrimitiveC {
+ public:
+  Roll() : PrimitiveC(kNameRoll) { InitIOName({"input_x"}, {"output"}); }
+  ~Roll() = default;
+  MS_DECLARE_PARENT(Roll, PrimitiveC);
+};
+
+AbstractBasePtr RollInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                          const std::vector<AbstractBasePtr> &input_args);
+using PrimRollPtr = std::shared_ptr<Roll>;
+}  // namespace ops
+}  // namespace mindspore
+#endif  // MINDSPORE_CORE_OPS_ROLL_H_

mindspore/nn/layer/basic.py

@@ -33,7 +33,7 @@ from ..cell import Cell
 from .activation import get_activation
 
 __all__ = ['Dropout', 'Flatten', 'Dense', 'ClipByNorm', 'Norm', 'OneHot', 'Pad', 'Unfold',
-           'Tril', 'Triu', 'ResizeBilinear', 'MatrixDiag', 'MatrixDiagPart', 'MatrixSetDiag', 'L1Regularizer']
+           'Tril', 'Triu', 'ResizeBilinear', 'MatrixDiag', 'MatrixDiagPart', 'MatrixSetDiag', 'L1Regularizer', 'Roll']
 
 
 class L1Regularizer(Cell):
@@ -1355,3 +1355,88 @@ class MatrixSetDiag(Cell):
         assist = _get_matrix_diag_part_assist(x_shape, x_dtype)
         out_matrix_set_diag = self.matrix_set_diag(input_x, diagonal, assist)
         return out_matrix_set_diag
+
+
+@constexpr
+def _check_input_dim(axis, dim, cls_name):
+    Validator.check_int_range(axis, -dim, dim, Rel.INC_LEFT, 'axis', cls_name)
+
+
+class Roll(Cell):
+    """
+    Rolls the elements of a tensor along an axis.
+
+    The elements are shifted positively (towards larger indices) by the offset of `shift` along the dimension of `axis`.
+    Negative `shift` values will shift elements in the opposite direction. Elements that roll passed the last position
+    will wrap around to the first and vice versa. Multiple shifts along multiple axes may be specified.
+
+    Args:
+        shift (Union[list(int), tuple(int), int]): Specifies the number of places by which elements are shifted
+            positively (towards larger indices) along the specified dimension. Negative shifts will roll the elements
+            in the opposite direction.
+        axis (Union[list(int), tuple(int), int]): Specifies the dimension indexes of shape to be rolled.
+
+    Inputs:
+        - **input_x** (Tensor) - Input tensor.
+
+    Outputs:
+        Tensor, has the same shape and type as `input_x`.
+
+    Raises:
+        TypeError: If `shift` is not an int, a tuple or a list.
+        TypeError: If `axis` is not an int, a tuple or a list.
+        TypeError: If element of `shift` is not an int.
+        TypeError: If element of `axis` is not an int.
+        ValueError: If axis is out of the range [-len(input_x.shape), len(input_x.shape)).
+        ValueError: If length of shape of `shift` is not equal to length of shape of `axis`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> input_x = Tensor(np.array([0, 1, 2, 3, 4]).astype(np.float32))
+        >>> op = nn.Roll(shift=2, axis=0)
+        >>> output = op(input_x)
+        >>> print(output)
+        [3. 4. 0. 1. 2.]
+        >>> input_x = Tensor(np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]).astype(np.float32))
+        >>> op = nn.Roll(shift=[1, -2], axis=[0, 1])
+        >>> output = op(input_x)
+        >>> print(output)
+        [[7. 8. 9. 5. 6.]
+         [2. 3. 4. 0. 1.]]
+    """
+
+    def __init__(self, shift, axis):
+        """Initialize Roll"""
+        super(Roll, self).__init__()
+        Validator.check_value_type("shift", shift, [int, tuple, list], self.cls_name)
+        Validator.check_value_type("axis", axis, [int, tuple, list], self.cls_name)
+        self.shape_op = P.Shape()
+        self.shift = shift
+        self.axis = axis
+        self.op_list = []
+
+        if not isinstance(self.axis, (list, tuple)):
+            self.op_list.append((inner.Roll(shift=self.shift, axis=0), self.axis))
+        else:
+            if len(self.shift) != len(self.axis):
+                raise ValueError('The shape of shift and the shape of axis must be the same.')
+            for idx, _ in enumerate(self.axis):
+                self.op_list.append((inner.Roll(shift=self.shift[idx], axis=0), self.axis[idx]))
+
+    def construct(self, input_x):
+        dim = len(self.shape_op(input_x))
+        for single_op_roll, single_axis in self.op_list:
+            _check_input_dim(single_axis, dim, self.cls_name)
+            if single_axis < 0:
+                single_axis += dim
+            transpose_perm = []
+            for i in range(dim):
+                transpose_perm.append(i)
+            transpose_perm[0], transpose_perm[single_axis] = single_axis, 0
+
+            input_x = input_x.transpose(transpose_perm)
+            input_x = single_op_roll(input_x)
+            input_x = input_x.transpose(transpose_perm)
+        return input_x

mindspore/ops/_grad_experimental/grad_inner_ops.py

@@ -31,3 +31,17 @@ def get_bprop_tensor_copy_slices(self):
         return x_grad, update_grad, zeros_like(begin), zeros_like(end), zeros_like(stride)
 
     return bprop
+
+
+@bprop_getters.register(inner.Roll)
+def get_bprop_roll(self):
+    """Generate bprop for Roll"""
+    shift = self.shift
+    axis = self.axis
+    roll_grad = inner.Roll(-shift, axis)
+
+    def bprop(x_input, out, dout):
+        dx = roll_grad(dout)
+        return (dx,)
+
+    return bprop

mindspore/ops/_op_impl/tbe/__init__.py

@@ -393,6 +393,7 @@ from .not_equal_ds import _not_ds_equal_tbe
 from .reciprocal_ds import _reciprocal_ds_tbe
 from .ctc_loss_v2 import _ctc_loss_v2_tbe
 from .ctc_loss_v2_grad import _ctc_loss_v2_grad_tbe
+from .roll import _roll_tbe
 from .soft_shrink import _soft_shrink_tbe
 from .soft_shrink_grad import _soft_shrink_grad_tbe
 from .hsigmoid_grad import _hsigmoid_grad_tbe

mindspore/ops/_op_impl/tbe/roll.py

@@ -0,0 +1,42 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+
+"""Roll op"""
+from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
+
+roll_op_info = TBERegOp("Roll") \
+    .fusion_type("OPAQUE") \
+    .async_flag(False) \
+    .binfile_name("roll.so") \
+    .compute_cost(10) \
+    .kernel_name("roll") \
+    .partial_flag(True) \
+    .attr("shift", "required", "listInt", "all") \
+    .attr("axis", "optional", "listInt", "all") \
+    .input(0, "input_x", False, "required", "all") \
+    .output(0, "output", False, "required", "all") \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.F16_Default) \
+    .dtype_format(DataType.I32_Default, DataType.I32_Default) \
+    .dtype_format(DataType.U32_Default, DataType.U32_Default) \
+    .dtype_format(DataType.I8_Default, DataType.I8_Default) \
+    .dtype_format(DataType.U8_Default, DataType.U8_Default) \
+    .get_op_info()
+
+
+@op_info_register(roll_op_info)
+def _roll_tbe():
+    """Roll TBE register"""
+    return

mindspore/ops/operations/_inner_ops.py

@@ -1227,3 +1227,69 @@ class TensorCopySlices(Primitive):
     def __init__(self):
         """Initialize TensorScatterUpdate"""
         self.init_prim_io_names(inputs=['x', 'value', 'begin', 'end', 'strides'], outputs=['y'])
+
+
+class Roll(Primitive):
+    """
+    Rolls the elements of a tensor along an axis.
+
+    The elements are shifted positively (towards larger indices) by the offset of `shift` along the dimension of `axis`.
+    Negative `shift` values will shift elements in the opposite direction. Elements that roll passed the last position
+    will wrap around to the first and vice versa. Multiple shifts along multiple axes may be specified.
+
+    Note:
+        This inner operation is valid only if the axis is equal to 0. If the shift and the axis are tuples or lists,
+        this inner operation is valid only for the first pair of elements.
+
+    Args:
+        shift (Union[list(int), tuple(int), int]): Specifies the number of places by which elements are shifted
+            positively (towards larger indices) along the specified dimension. Negative shifts will roll the elements
+            in the opposite direction.
+        axis (Union[list(int), tuple(int), int]): Specifies the dimension indexes of shape to be rolled. The value is
+            forced to be zero in this operation.
+
+    Inputs:
+        - **input_x** (Tensor) - Input tensor.
+
+    Outputs:
+        Tensor, has the same shape and type as `input_x`.
+
+    Raises:
+        TypeError: If `shift` is not an int, a tuple or a list.
+        TypeError: If `axis` is not an int, a tuple or a list.
+        TypeError: If element of `shift` is not an int.
+        TypeError: If element of `axis` is not an int.
+        ValueError: If axis is not equal to 0.
+        ValueError: If shape of `shift` is not equal to 1.
+        ValueError: If shape of `axis` is not equal to 1.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.ops.operations import _inner_ops as inner
+        >>> input_x = Tensor(np.array([0, 1, 2, 3, 4]).astype(np.float32))
+        >>> op = inner.Roll(shift=2, axis=0)
+        >>> output = op(input_x)
+        >>> print(output)
+        [3. 4. 0. 1. 2.]
+        >>> input_x = Tensor(np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]).astype(np.float32))
+        >>> op = inner.Roll(shift=-1, axis=0)
+        >>> output = op(input_x)
+        >>> print(output)
+        [[5. 6. 7. 8. 9.]
+         [0. 1. 2. 3. 4.]]
+    """
+
+    @prim_attr_register
+    def __init__(self, shift, axis):
+        """Initialize Roll"""
+        validator.check_value_type("shift", shift, [int, tuple, list], self.name)
+        validator.check_value_type("axis", axis, [int, tuple, list], self.name)
+        if isinstance(shift, (tuple, list)) and isinstance(axis, (tuple, list)):
+            validator.check_equal_int(len(shift), 1, "shift size", self.name)
+            validator.check_equal_int(len(axis), 1, "shift size", self.name)
+            validator.check_equal_int(axis[0], 0, "axis", self.name)
+        elif isinstance(shift, int) and isinstance(axis, int):
+            validator.check_equal_int(axis, 0, "axis", self.name)
+        self.init_prim_io_names(inputs=['input_x'], outputs=['output'])

tests/ut/python/ops/test_ops.py

@@ -2189,6 +2189,10 @@ test_case_nn_ops = [
                         Tensor(np.zeros((1, 1, 2, 2)), mstype.uint16)],
         'desc_bprop': [],
         'skip': ['backward']}),
+    ('Roll', {
+        'block': nn.Roll(shift=[1, -2], axis=[0, 1]),
+        'desc_inputs': [Tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], mstype.float32)],
+        'desc_bprop': [Tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], mstype.float32)]}),
     ('SoftShrink', {
         'block': P.SoftShrink(),
         'desc_inputs': [Tensor(np.array([[0.5297, 0.7871, 1.1754], [0.7836, 0.6218, -1.1542]]), mstype.float32)],