mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/ops/operations/_inner_ops.py

@@ -19,22 +19,24 @@ from collections.abc import Iterable
 import numpy as np
 
 from mindspore.common import Tensor
+from mindspore.common._stub_tensor import StubTensor
 from mindspore.ops import composite as C
 from mindspore.ops.operations.array_ops import Cast
 from mindspore.ops.operations._scalar_ops import bit_or, bit_and
+from mindspore.ops.operations.comm_ops import ReduceOp
 from mindspore.ops import signature as sig
 from mindspore.ops.operations.math_ops import _infer_shape_reduce
-from mindspore.ops.primitive import PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register, Primitive, _run_op
-from mindspore import context
+from mindspore.ops.primitive import PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register, Primitive,\
+    _run_op, _check_contains_variable
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore._c_expression import typing
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
-from mindspore.communication.management import GlobalComm
+from mindspore.communication.management import GlobalComm, get_rank
 from mindspore.common.api import _pynative_executor
 from mindspore.common._register_for_adapter import ms_adapter_registry
-
+from mindspore import ops
 
 # Bit operation
 bit_and = bit_and()
@@ -73,12 +75,11 @@ class ExtractImagePatches(Primitive):
             - valid: Means that the taken patch area must be completely covered in the original image.
 
     Inputs:
-        - **input_x** (Tensor) - A 4-D tensor whose shape is [in_batch, in_depth, in_row, in_col] and
-          data type is number.
+        - **input_x** (Tensor) - A 4-D tensor whose shape is :math:`(in\_batch, in\_depth, in\_row, in\_col)`.
 
     Outputs:
-        Tensor, a 4-D tensor whose data type is same as 'input_x',
-        and the shape is [out_batch, out_depth, out_row, out_col], Where the out_batch is the same as the in_batch
+        Tensor, a 4-D tensor whose data type is same as 'input_x', and the shape
+        is :math:`(out\_batch, out\_depth, out\_row, out\_col)`,where the out_batch is the same as the in_batch
         and
 
         .. math::
@@ -121,7 +122,6 @@ class ExtractImagePatches(Primitive):
         validator.check_value_type('padding', padding, [str], self.name)
         self.padding = validator.check_string(padding.upper(), ['VALID', 'SAME'], 'padding', self.name)
         self.add_prim_attr("padding", self.padding)
-        self.is_ge = context.get_context("enable_ge")
 
 
 class Quant(PrimitiveWithInfer):
@@ -144,7 +144,7 @@ class Quant(PrimitiveWithInfer):
     Args:
         scale (float) : Specifies the scaling ratio.
         offset (float): Specifies the offset.
-        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False.
+        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: ``False``.
         round_mode (str): Specifies the way to round. Must be one of ["Round", "Floor", "Ceil", "Trunc"].
           Default: "Round".
 
@@ -172,7 +172,7 @@ class Quant(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_type):
-        validator.check_subclass("input_x", x_type, mstype.tensor, self.name)
+        validator.check_subclass("input_x", x_type, mstype.tensor_type, self.name)
         validator.check_type_name("input_x", x_type, [mstype.float16, mstype.float32], self.name)
         return mstype.int8
 
@@ -254,8 +254,8 @@ class Dequant(PrimitiveWithInfer):
         This operation only support Ascend 310 inference environment.
 
     Args:
-        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False.
-        relu_flag (bool): Specifies whether to perform ReLU. Default: False.
+        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: ``False``.
+        relu_flag (bool): Specifies whether to perform ReLU. Default: ``False``.
 
     Inputs:
         - **input_x** (Tensor) : Input tensor. Must be mindspore.int32.
@@ -281,7 +281,7 @@ class Dequant(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_type, deq_scale_type):
-        validator.check_subclass("x", x_type, mstype.tensor, self.name)
+        validator.check_subclass("x", x_type, mstype.tensor_type, self.name)
         validator.check_type_name("x", x_type, [mstype.int32], self.name)
         validator.check_type_name("deq_scale", deq_scale_type, [mstype.float16, mstype.uint64], self.name)
         return mstype.float16
@@ -502,6 +502,109 @@ class Receive(PrimitiveWithInfer):
         return self.get_attr_dict()['dtype']
 
 
+class Reduce(PrimitiveWithInfer):
+    """
+    Reduces tensor across the processes in the specified communication group.
+
+    Note:
+        Only process with destination rank receives the reduced output.
+        Other processes only get a tensor with shape [1], which has no mathematical meaning.
+
+    Args:
+        dest_rank (int): Specifies the rank of the process that receives the reduced output.
+        op (str, optional): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
+                On the CPU, only 'sum' is supported. Default: ``ReduceOp.SUM`` .
+        group (str, optional): The communication group to work on.
+            Default: "hccl_world_group" on Ascend, "nccl_world_group" on GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> import numpy as np
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class ReduceNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.reduce = ops.Reduce(dest_rank=1)
+        >>>
+        >>>     def construct(self, x):
+        >>>         out = self.reduce(x)
+        >>>         return out
+        >>> input = Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> net = ReduceNet()
+        >>> output = net(input)
+        >>> print(output)
+        Process with rank 1: [[4. 4. 4. 4. 4. 4. 4. 4.]
+                             [4. 4. 4. 4. 4. 4. 4. 4.]],
+        Other proesses: [0.].
+    """
+
+    @prim_attr_register
+    def __init__(self, dest_rank, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
+        self.dest_rank = dest_rank
+        self.op = op
+        self.group = group
+
+    def infer_shape(self, x_shape):
+        # The process with dest_rank returns the reduced output.
+        # Other processes only gets a tensor with shape [1], which has no mathematical meaning.
+        if self.dest_rank == get_rank():
+            return x_shape
+        return [1]
+
+    def infer_dtype(self, x_dtype):
+        return x_dtype
+
+
+class Barrier(PrimitiveWithInfer):
+    """
+    Synchronizes all processes in the specified group.
+
+    Note:
+        After calling this collective operator,
+        this process will be blocked until all other processes in the group call this operator.
+
+    Args:
+        group (str, optional): The communication group to work on.
+            Default: "hccl_world_group" on Ascend, "nccl_world_group" on GPU.
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> import numpy as np
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class BarrierNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.barrier = ops.Barrier()
+        >>>
+        >>>     def construct(self):
+        >>>         self.barrier()
+        >>> net = BarrierNet()
+        >>> net()
+    """
+
+    @prim_attr_register
+    def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
+        self.group = group
+        self.add_prim_attr("side_effect_mem", True)
+
+    def infer_shape(self):
+        return [1]
+
+    def infer_dtype(self):
+        return mstype.float32
+
+
 class MatrixSetDiag(PrimitiveWithInfer):
     r"""
     Modifies the batched diagonal part of a batched tensor.
@@ -604,9 +707,9 @@ class ConfusionMulGrad(PrimitiveWithInfer):
         return outshape0, outshape1
 
     def infer_dtype(self, input0_dtype, input1_dtype, input2_dtype):
-        validator.check_subclass("input0_dtype", input0_dtype, mstype.tensor, self.name)
-        validator.check_subclass("input1_dtype", input1_dtype, mstype.tensor, self.name)
-        validator.check_subclass("input2_dtype", input2_dtype, mstype.tensor, self.name)
+        validator.check_subclass("input0_dtype", input0_dtype, mstype.tensor_type, self.name)
+        validator.check_subclass("input1_dtype", input1_dtype, mstype.tensor_type, self.name)
+        validator.check_subclass("input2_dtype", input2_dtype, mstype.tensor_type, self.name)
         return input0_dtype, input1_dtype
 
 
@@ -619,7 +722,7 @@ class ConvertToDynamic(PrimitiveWithCheck):
 
     Args:
         is_dynamic_rank (bool): If true, convert to dynamic rank.
-                                If false, convert to dynamic shape. Default: False.
+                                If false, convert to dynamic shape. Default: ``False``.
 
     Inputs:
         - **input** (Tensor) - The tensor used for testing.
@@ -664,7 +767,7 @@ class ConvertToDynamic(PrimitiveWithCheck):
         validator.check("input_shape rank", len(input_shape), "", 0, validator.GT, self.name)
 
     def check_dtype(self, input_dtype):
-        validator.check_subclass("input_dtype", input_dtype, mstype.tensor, self.name)
+        validator.check_subclass("input_dtype", input_dtype, mstype.tensor_type, self.name)
 
 
 class GpuConvertToDynamicShape(PrimitiveWithCheck):
@@ -714,7 +817,7 @@ class GpuConvertToDynamicShape(PrimitiveWithCheck):
         validator.check("input_shape rank", len(input_shape), "", 0, validator.GT, self.name)
 
     def check_dtype(self, input_dtype):
-        validator.check_subclass("input_dtype", input_dtype, mstype.tensor, self.name)
+        validator.check_subclass("input_dtype", input_dtype, mstype.tensor_type, self.name)
 
 
 class ErrorOnDynamicShapeInput(PrimitiveWithInfer):
@@ -766,7 +869,7 @@ class ErrorOnDynamicShapeInput(PrimitiveWithInfer):
 
     def infer_type(self, input_dtype):
         """Infer the dtype of input for ErrorOnDynamicShapeInput."""
-        validator.check_subclass("input_dtype", input_dtype, mstype.tensor, self.name)
+        validator.check_subclass("input_dtype", input_dtype, mstype.tensor_type, self.name)
         return input_dtype
 
     def infer_value(self, input_tensor):
@@ -816,7 +919,7 @@ class SequenceMask(PrimitiveWithCheck):
         validator.check("maxlen_shape", len(maxlen_shape), "", 0, validator.EQ, self.name)
 
     def check_dtype(self, lengths_dtype, maxlen_dtype):
-        validator.check_subclass("lengths_dtype", lengths_dtype, mstype.tensor, self.name)
+        validator.check_subclass("lengths_dtype", lengths_dtype, mstype.tensor_type, self.name)
         validator.check_subclass("maxlen", maxlen_dtype, mstype.number, self.name)
 
 
@@ -1170,8 +1273,8 @@ class DynamicStitch(PrimitiveWithCheck):
         return out_shape
 
     def check_dtype(self, indices_type, data_type):
-        validator.check_subclass("indices[0]", indices_type[0], mstype.tensor, self.name)
-        validator.check_subclass("data[0]", data_type[0], mstype.tensor, self.name)
+        validator.check_subclass("indices[0]", indices_type[0], mstype.tensor_type, self.name)
+        validator.check_subclass("data[0]", data_type[0], mstype.tensor_type, self.name)
         indices_num = len(indices_type)
         for i in range(0, indices_num):
             validator.check_tensor_dtype_valid(f'indices[{i}]', indices_type[i], mstype.int32, self.name)
@@ -1418,6 +1521,7 @@ class DecodeImage(PrimitiveWithInfer):
 
     Examples:
     """
+
     @prim_attr_register
     def __init__(self, channels=0, dtype=mstype.uint8, expand_animations=False, _op_max_shape="8192,8192,3",
                  _op_max_size=[8000000]):
@@ -1467,7 +1571,7 @@ class DynamicBroadcastTo(Primitive):
     Inputs:
         - **input_x** (Tensor) - The input tensor. The data type should be one of the following types:
           float16, float32, int32, int8, uint8.
-          The shape is :math:`(N,*)` where :math:`*` means,any number of additional dimensions.
+          The shape is :math:`(N,*)` where :math:`*` means any number of additional dimensions.
         - **shape** (Tensor): The target shape to broadcast.
 
     Outputs:
@@ -1495,6 +1599,16 @@ class Cummin(Primitive):
 
     Refer to :func:`mindspore.ops.cummin` for more detail.
 
+    Args:
+        axis (int): The axis to accumulate the tensor's value. Must be in the range [-rank(input), rank(input)).
+
+    Inputs:
+        - **input** (Tensor) - The input tensor.
+
+    Outputs:
+        A tuple of 2 Tensors(values, indices), containing the cumulative minimum of elements and the index,
+        The shape of each output tensor is the same as input `input`.
+
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
@@ -1509,6 +1623,7 @@ class Cummin(Primitive):
         >>> print(output[1])
         [0 1 1 1 4 4]
     """
+
     @prim_attr_register
     def __init__(self, axis):
         """Initialize Cummin"""
@@ -1528,7 +1643,7 @@ class DynamicResizeNearestNeighbor(Primitive):
 
     Args:
         align_corners (bool): Whether the centers of the 4 corner pixels of the input
-                              and output tensors are aligned. Default: False.
+                              and output tensors are aligned. Default: ``False``.
 
     Inputs:
         - **input_x** (Tensor) - The input tensor. The shape of the tensor is :math:`(N, C, H, W)`.
@@ -1613,7 +1728,7 @@ class PsROIPooling(PrimitiveWithInfer):
         return output_shape, output_map_shape
 
     def infer_dtype(self, inputs_type, rois_type):
-        map_type = mstype.tensor_type(mstype.int32)
+        map_type = mstype.TensorType(mstype.int32)
         return inputs_type, map_type
 
 
@@ -1671,8 +1786,10 @@ class PartitionedCall(PrimitiveWithInfer):
 
     Examples:
     """
+
     @prim_attr_register
     def __init__(self, graph, executor_type=""):
+        super(PartitionedCall, self).__init__(self.__class__.__name__)
         self.add_prim_attr("executor_type", executor_type)
         self.graph = graph
 
@@ -1744,9 +1861,6 @@ class CellBackwardHook(PrimitiveWithInfer):
     def __call__(self, args):
         if not isinstance(args, tuple):
             args = (args,)
-        for arg in args:
-            if isinstance(arg, Parameter) and arg.has_init:
-                arg.init_data()
         return _run_op(self, self.name, args)
 
     def infer_shape(self, *inputs_shape):
@@ -1832,16 +1946,32 @@ class Format(PrimitiveWithInfer):
     def __init__(self):
         self.init_prim_io_names(inputs=['string', 'args'], outputs=['string'])
 
+
     def __infer__(self, str_, *var):
-        str_value = str_["value"]
+        def check_variable(str_, var):
+            if _check_contains_variable(str_['dtype'], str_['value']):
+                return True
+
+            for item in var:
+                if _check_contains_variable(item['dtype'], item['value']):
+                    return True
+            return False
+
+
+        if check_variable(str_, var):
+            return {'dtype': mstype.string, 'shape': [], 'value': None}
+
+
+        str_value = str_['value']
+        kwargs = dict()
         var_value = list()
-        if str_value is None and str_["dtype"] is not None:
-            raise ValueError("str.format not support to input a variable.")
+
         for item in var:
-            if item["value"] is None and item["dtype"] is not None:
-                raise ValueError("str.format not support to input a variable.")
+            if isinstance(item["dtype"], typing.Keyword):
+                kwargs.update(item["value"])
             var_value.append(item["value"])
-        value = str_value.format(*var_value)
+
+        value = str_value.format(*var_value, **kwargs)
         return {'dtype': mstype.string, 'shape': [], 'value': value}
 
 
@@ -1982,7 +2112,7 @@ class ClipByNorm(PrimitiveWithInfer):
 
     Args:
         axis (Union[None, int, tuple(int), list(int)]): Compute the `L_2`-norm along the specific dimension.
-                                                       Default: None, all dimensions to calculate.
+                                                       Default: ``None``, all dimensions to calculate.
 
     Inputs:
         - **x** (Tensor) - Tensor of shape N-D. The type must be float16 or float32.
@@ -2060,8 +2190,8 @@ class TopTypeof(Primitive):
             'slice': mstype.Slice(),
             'list': mstype.List(),
             'tuple': mstype.Tuple(),
-            'Tensor': mstype.tensor,
-            'NoneType': mstype.none_type(),
+            'Tensor': mstype.tensor_type,
+            'NoneType': mstype.NoneType(),
             'int': mstype.Int(),
             'bool': mstype.Bool(),
             'ellipsis': mstype.Ellipsis_(),
@@ -2098,7 +2228,7 @@ class MixedPrecisionCast(Primitive):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor
-        >>> from mindspore.common import dtype as mstype
+        >>> from mindspore import dtype as mstype
         >>> from mindspore.ops.operations import _inner_ops as inner
         >>> x = Tensor(np.ones([2, 3], dtype=np.float32))
         >>> out = inner.MixedPrecisionCast(mstype.float16, x)
@@ -2175,13 +2305,22 @@ class CheckBprop(PrimitiveWithInfer):
             raise ValueError(f"For {tips} the number of return values(gradients) must be equal to "
                              f"the number of input arguments except 'out' and 'dout', "
                              f"which is:{len(yshapes)} but got {len(xshapes)}.")
-        checking_range = len(yshapes)
-        for i in range(checking_range):
-            xshape = xshapes[i]
-            yshape = yshapes[i]
+
+        def shape_equal(shape1, shape2):
+            if len(shape1) != len(shape2):
+                return False
+            for shape_axis1, shape_axis2 in zip(shape1, shape2):
+                if shape_axis1 == -1 or shape_axis2 == -1:
+                    continue
+                if shape_axis1 != shape_axis2:
+                    return False
+            return True
+
+        for i, (xshape, yshape) in enumerate(zip(xshapes, yshapes)):
             if not xshape or not yshape:
                 continue
-            if xshape != yshape:
+
+            if not shape_equal(xshape, yshape):
                 raise ValueError(f"For {tips}, the {i}th return value(gradient of the {i}th argument) "
                                  f"should have the same shape as the {i}th argument, "
                                  f"which is:{yshape}, but got: {xshape}.")
@@ -2200,18 +2339,19 @@ class CheckBprop(PrimitiveWithInfer):
         for i in range(checking_range):
             xdtype = xdtypes[i]
             ydtype = ydtypes[i]
-            if isinstance(xdtype, mstype.anything_type) or isinstance(ydtype, mstype.anything_type):
+            if isinstance(xdtype, mstype.AnythingType) or isinstance(ydtype, mstype.AnythingType):
                 continue
-            if isinstance(ydtype, mstype.function_type):
-                if not isinstance(xdtype, mstype.env_type_type):
+            if isinstance(ydtype, mstype.FunctionType):
+                if not isinstance(xdtype, mstype.EnvType):
                     raise TypeError(f"For {tips}, the {i}th return value(gradient of the {i}th argument) type "
-                                    f"should be {mstype.env_type_type}, but got {xdtype}.")
+                                    f"should be {mstype.EnvType}, but got {xdtype}.")
             if xdtype != ydtype:
                 raise TypeError(f"For {tips}, the {i}th return value(gradient of the {i}th argument) "
                                 f"should have the same dtype as the {i}th argument, "
                                 f"which is:{ydtype}, but got: {xdtype}.")
         return xdtypes
 
+
 check_bprop = CheckBprop()
 
 
@@ -2246,8 +2386,8 @@ class SameTypeShape(PrimitiveWithInfer):
         return x
 
     def __infer__(self, x, y):
-        validator.check_subclass('x', x['dtype'], mstype.tensor, self.name)
-        validator.check_subclass('y', y['dtype'], mstype.tensor, self.name)
+        validator.check_subclass('x', x['dtype'], mstype.tensor_type, self.name)
+        validator.check_subclass('y', y['dtype'], mstype.tensor_type, self.name)
         validator.check('x dtype', x['dtype'], 'y dtype', y['dtype'], validator.EQ, self.name, TypeError)
         validator.check('x shape', x['shape'], 'y shape', y['shape'], validator.EQ, self.name)
         return x
@@ -2374,13 +2514,15 @@ class ConvertToAdapterTensor(Primitive):
         >>> print(x)
         [1 2 3]
     """
+
     @prim_attr_register
     def __init__(self):
         """Initialize"""
 
     def __call__(self, x):
-        """run in PyNative mode"""
-        return ms_adapter_registry.tensor(x, inner=True)
+        """Run in PyNative mode"""
+        return ms_adapter_registry.tensor(x, cast_tensor=True)
+
 
 convert_to_adapter_tensor = ConvertToAdapterTensor()
 
@@ -2405,13 +2547,17 @@ class ConvertToMsTensor(Primitive):
         >>> print(x)
         [1 2 3]
     """
+
     @prim_attr_register
     def __init__(self):
         """Initialize"""
 
     def __call__(self, x):
-        """run in PyNative mode"""
-        return Tensor(x)
+        """Run in PyNative mode"""
+        if isinstance(x, StubTensor):
+            return StubTensor(stub=x.stub, tensor=x.tensor)
+        return ops.deepcopy(x)
+
 
 convert_to_ms_tensor = ConvertToMsTensor()
 
@@ -2458,6 +2604,7 @@ class IsParameter(PrimitiveWithInfer):
     """
         Check if input is `Parameter`
     """
+
     @prim_attr_register
     def __init__(self):
         """Initialize IsParameter"""
@@ -2468,7 +2615,7 @@ class IsParameter(PrimitiveWithInfer):
     def __infer__(self, x):
         return {'shape': [],
                 'dtype': mstype.bool_,
-                'value': isinstance(x['dtype'], mstype.ref_type)}
+                'value': isinstance(x['dtype'], mstype.RefType)}
 
 
 class SiLU(Primitive):
@@ -2547,3 +2694,98 @@ class SetitemTensorIndexInfo(Primitive):
 
     def __call__(self, data, index, value):
         return Tensor_.setitem_index_info(data, index, value, self.is_ascend)
+
+
+class IsConstant(Primitive):
+    r"""
+        Check if the input is constant
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize IsConstant"""
+
+    def __call__(self, x):
+        return True
+
+
+class SelectView(Primitive):
+    r"""
+        Select tensor of view
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        self.init_prim_io_names(inputs=['input_tensor', 'input_indices', 'axis'], outputs=['output'])
+
+
+class CopyWithSlice(Primitive):
+    r"""
+        Copy data to discontinuous tensor
+    """
+    @prim_attr_register
+    def __init__(self):
+        self.add_prim_attr('side_effect_mem', True)
+        self.init_prim_io_names(inputs=['x', 'y'], outputs=['x'])
+
+
+class MoeFFN(Primitive):
+    r"""
+    The MoeFFN computation is similar to Feed-Forward Network, it contains matmul + gelu + matmul.
+
+    Args:
+        activation (string): The activation type, set to 'fastgelu' or 'gelu'.
+        Only support 'fastgelu' for now. Default: "fastgelu".
+
+    Inputs:
+        - **x** (Tensor) - The input tensor with data type of int8, float16.
+          Input tensor of shape :math:`(batch\_size * seq\_length, hidden\_size)`.
+        - **expert_tokens** (Tensor]) - The expert tokens tensor with data type of int64.
+          Expert tokens tensor of shape :math:`(16,)`. For example, `(2, 1, 0, .., 9)`
+          indicate that the 0th expert deals with 2 tokens, the 1th expert deals with 1 tokens,
+          the 2th expert do noting and so on.
+        - **weight1** (Tensor) - The weight1 tensor with data type of float16.
+          Weight1 tensor of shape :math:`(expert\_num, hidden\_size, ffn\_hidden\_size)`.
+        - **bias1** (Tensor) - The bias1 tensor with data type of float16.
+          Bias1 tensor of shape :math:`(expert\_num, ffn\_hidden\_size)`.
+        - **weight2** (Tensor) - The weight2 tensor with data type of float16.
+          Weight2 tensor of shape :math:`(expert\_num, ffn\_hidden\_size, hidden\_size)`.
+        - **bias2** (Tensor) - The bias2 tensor with data type of float16.
+          Bias2 tensor of shape :math:`(expert\_num, hidden\_size)`.
+        - **scale** (Tensor) - The scale tensor with data type of float16. Not enable now.
+        - **offset** (Tensor) - The offset tensor with data type of float16. Not enable now.
+        - **deq_scale1** (Tensor) - The deq_scale1 tensor with data type of float16. Not enable now.
+        - **deq_scale2** (Tensor) - The deq_scale2 tensor with data type of float16. Not enable now.
+
+    Outputs:
+        Tensor of shape :math:`(batch\_size * seq\_length, hidden\_size)`. With data type of float16.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.ops.operations import _inner_ops
+        >>> b = 4
+        >>> s = 128
+        >>> h = 1024
+        >>> h_f = 4 * h
+        >>> e = 16
+        >>> x = Tensor(np.random.randn(b * s, h).astype(np.float16))
+        >>> expert_tokens = Tensor(np.random.randn(e).astype(np.int64))
+        >>> w1 = Tensor(np.random.randn(e, h, h_f).astype(np.float16))
+        >>> bias1 = Tensor(np.random.randn(e, h_f).astype(np.float16))
+        >>> w2 = Tensor(np.random.randn(e, h_f, h).astype(np.float16))
+        >>> bias2 = Tensor(np.random.randn(e, h).astype(np.float16))
+        >>> moe_ffn = _inner_ops.MoeFFN("fastgelu")
+        >>> output = moe_ffn(x, w1, bias1, w2, bias2)
+        >>> print(output)
+    """
+
+    @prim_attr_register
+    def __init__(self, activation):
+        """Initialize MoeFFN."""
+        self.init_prim_io_names(inputs=["x", "expert_tokens", "weight1", "bias1",
+                                        "weight2", "bias2", "scale", "offset", "deq_scale1"
+                                        "deq_scale2"],
+                                outputs=["y"])
+        self.activation = activation

mindspore/ops/operations/_ms_kernel.py

@@ -486,9 +486,9 @@ def kernel(fn=None, reg_info=None, compile_attrs=None):
     will enjoy the automatic dtype/shape infer for free.
 
     Args:
-        fn (Function): The Python function that will be run as a custom operator. Default: None.
-        reg_info (tuple[str, dict]): Each item represents registration information in json format. Default: None.
-        compile_attrs (Dict): The Python object is used to distinguish the compiled function. Default: None.
+        fn (Function): The Python function that will be run as a custom operator. Default: ``None`` .
+        reg_info (tuple[str, dict]): Each item represents registration information in json format. Default: ``None`` .
+        compile_attrs (Dict): The Python object is used to distinguish the compiled function. Default: ``None`` .
 
     Returns:
         Function, if `fn` is not None, returns a callable function that will execute the Hybrid DSL function;