mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/layer/dense.py

@@ -77,6 +77,7 @@ class BiDense(Cell):
         bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter.
             The values of str refer to the function `initializer`. Default: ``None`` .
         has_bias (bool): Specifies whether the layer uses :math:`\text{bias}` vector. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Shape:
         - **input1** - :math:`(*, H_{in1})` where :math:`H_{in1}=\text{in1_channels}` and
@@ -90,8 +91,8 @@ class BiDense(Cell):
           are the same shape as the inputs.
 
     Dtype:
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2**.
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input1**.
+        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2** .
+        - **input2** (Tensor) - The dtype must be float16 or float32 and be same as **input1** .
         - **output** (Tensor) - With the same dtype as the inputs.
 
     Weights:
@@ -133,7 +134,8 @@ class BiDense(Cell):
                  out_channels,
                  weight_init=None,
                  bias_init=None,
-                 has_bias=True):
+                 has_bias=True,
+                 dtype=mstype.float32):
         super().__init__()
         self.in_channels = Validator.check_positive_int(in1_channels, "in1_channels", self.cls_name)
         self.in_channels = Validator.check_positive_int(in2_channels, "in2_channels", self.cls_name)
@@ -156,7 +158,8 @@ class BiDense(Cell):
                                  f"equal to 'in2_channels'. But got 'weight_init': {weight_init}, "
                                  f"'out_channels': {out_channels}, 'in_channels': {in1_channels}, "
                                  f"'in2_channels': {in2_channels}")
-        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels)), 'weight')
+        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels), dtype=dtype),
+                                'weight')
 
         if self.has_bias:
             if bias_init is None:
@@ -166,7 +169,7 @@ class BiDense(Cell):
                     raise ValueError(f"For '{self.cls_name}', bias init shape error. The ndim of 'bias_init' should "
                                      f"be equal to 1, and the first dim must be equal to 'out_channels'. But got "
                                      f"'bias_init': {bias_init}, 'out_channels': {out_channels}.")
-            self.bias = Parameter(initializer(bias_init, [out_channels]), name="bias")
+            self.bias = Parameter(initializer(bias_init, [out_channels], dtype=dtype), name="bias")
             self.bias_add = P.BiasAdd()
         self.matmul = P.MatMul()
 

mindspore/nn/layer/embedding.py

@@ -64,11 +64,13 @@ class Embedding(Cell):
         embedding_size (int): The size of each embedding vector.
         use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: ``False`` .
         embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
-            Refer to class `initializer` for the values of string when a string
-            is specified. Default: ``'normal'`` .
+            Refer to class `mindspore.common.initializer
+            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            for the values of string when a string is specified. Default: ``'normal'`` .
         dtype (:class:`mindspore.dtype`): Data type of `x`. Default: ``mstype.float32`` .
         padding_idx (int, None): When the padding_idx encounters index, the output embedding vector of this index
                                  will be initialized to zero. Default: ``None`` . The feature is inactivated.
+
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(\text{batch_size}, \text{x_length})`. The elements of
           the Tensor must be integer and not larger than vocab_size. Otherwise the corresponding embedding vector will
@@ -145,9 +147,8 @@ class Embedding(Cell):
         return output
 
     def extend_repr(self):
-        s = 'vocab_size={}, embedding_size={}, use_one_hot={}, embedding_table={}, dtype={}, padding_idx={}'.format(
-            self.vocab_size, self.embedding_size, self.use_one_hot, self.embedding_table, self.dtype, self.padding_idx)
-        return s
+        return f'vocab_size={self.vocab_size}, embedding_size={self.embedding_size}, use_one_hot={self.use_one_hot}, ' \
+            f'embedding_table={self.embedding_table}, dtype={self.dtype}, padding_idx={self.padding_idx}'
 
 
 @_primexpr
@@ -190,6 +191,7 @@ class EmbeddingLookup(Cell):
             parameter server trainning mode and 'DEVICE' target. And the moment parameter of corresponding
             optimizer will also be set to the cache size. In addition, it should be noted that it will cost the 'DEVICE'
             memory, so suggests setting a reasonable value to avoid insufficient memory.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(y_1, y_2, ..., y_S)`.
@@ -231,7 +233,7 @@ class EmbeddingLookup(Cell):
 
     def __init__(self, vocab_size, embedding_size, param_init='normal',
                  target='CPU', slice_mode='batch_slice', manual_shapes=None,
-                 max_norm=None, sparse=True, vocab_cache_size=0):
+                 max_norm=None, sparse=True, vocab_cache_size=0, dtype=mstype.float32):
         """Initialize EmbeddingLookup."""
         super(EmbeddingLookup, self).__init__()
         Validator.check_value_type('sparse', sparse, [bool], self.cls_name)
@@ -255,8 +257,8 @@ class EmbeddingLookup(Cell):
         if enable_ps:
             self._process_vocab_cache(slice_mode)
         self.embedding_size = Validator.check_positive_int(embedding_size, 'embedding_size', self.cls_name)
-        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size]),
-                                         name='embedding_table')
+        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size],
+                                                     dtype=dtype), name='embedding_table')
         parallel_mode = _get_parallel_mode()
         is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
         self.gather_revert = P.Gather()
@@ -267,7 +269,7 @@ class EmbeddingLookup(Cell):
         if is_auto_parallel:
             self.unique = P.Unique().shard(((1,),))
         if self.cache_enable and enable_ps:
-            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init)
+            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init, dtype=dtype)
             if is_auto_parallel:
                 self.unique.add_prim_attr('cache_enable', True)
         indices_shape_size = 2
@@ -310,8 +312,8 @@ class EmbeddingLookup(Cell):
         else:
             if is_auto_parallel:
                 support_mode = ["field_slice", "table_row_slice", "table_column_slice", "batch_slice"]
-                raise ValueError("For '{}', the 'slice_mode' must be in {}, "
-                                 "but got \"{}\".".format(self.cls_name, support_mode, slice_mode))
+                raise ValueError(f"For '{self.cls_name}', the 'slice_mode' must be in {support_mode}, "
+                                 f"but got \"{slice_mode}\".")
         if self.cache_enable and not enable_ps:
             raise ValueError(f"For '{self.cls_name}', haven't supported cache enable for not ps mode.")
         self.embedding_table.unique = self.forward_unique
@@ -354,7 +356,8 @@ class EmbeddingLookup(Cell):
             if _is_role_worker():
                 self.vocab_size = self.vocab_cache_size
 
-    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init):
+    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init,
+                                       dtype=mstype.float32):
         """PS embeddingLookup cache enable set."""
         if self.sparse:
             self.forward_unique = True
@@ -368,10 +371,10 @@ class EmbeddingLookup(Cell):
         if _enable_distributed_mindrt():
             self.rank_id = get_rank()
             if self.is_ps_server:
-                self._slice_pserver_embeddings("zeros")
+                self._slice_pserver_embeddings("zeros", dtype=dtype)
                 self._set_cache_enable_and_key_for_pserver(param_key)
 
-    def _slice_pserver_embeddings(self, param_init):
+    def _slice_pserver_embeddings(self, param_init, dtype=mstype.float32):
         '''
         Method to slice embedding tables on Parameter Servers.
         It helps to train with a large scale embedding table and is used only in Parameter Server training mode.
@@ -399,7 +402,7 @@ class EmbeddingLookup(Cell):
         for i in range(server_num):
             self.embedding_table_list.append(Parameter(initializer(param_init,
                                                                    [self.embedding_table_vocab_dim_list[i],
-                                                                    self.embedding_size]),
+                                                                    self.embedding_size], dtype=dtype),
                                                        name="embedding_table_server_" + str(i)))
 
             self.embedding_offset.append(offset)
@@ -505,12 +508,13 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
             :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'``.
         feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently.
             Default:  ``None`` .
-        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32
-                                       or None. Default: ``None`` .
+        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32.
+            Default: ``None`` .
         sparse (bool): Using sparse mode. When 'target' is set to ``'CPU'`` , 'sparse' has to be true.
             Default: ``True`` .
         operator (str): The pooling method for the features in one field. Support ``'SUM'`` , ``'MEAN'`` and
             ``'MAX'`` . Default: ``'SUM'`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(batch\_size, seq\_length)`.
@@ -529,12 +533,12 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         TypeError: If `vocab_size` or `embedding_size` or `field_size` is not an int.
         TypeError: If `sparse` is not a bool or `feature_num_list` is not a tuple.
         ValueError: If `vocab_size` or `embedding_size` or `field_size` is less than 1.
-        ValueError: If `target` is neither 'CPU' nor 'DEVICE'.
-        ValueError: If `slice_mode` is not one of 'batch_slice', 'field_slice', 'table_row_slice',
-                    'table_column_slice'.
-        ValueError: If `sparse` is False and `target` is 'CPU'.
-        ValueError: If `slice_mode` is 'field_slice' and `feature_num_list` is None.
-        ValueError: If `operator` is not one of 'SUM', 'MAX', 'MEAN'.
+        ValueError: If `target` is neither ``'CPU'`` nor ``'DEVICE'``.
+        ValueError: If `slice_mode` is not one of ``'batch_slice'``, ``'field_slice'``, ``'table_row_slice'``,
+                    ``'table_column_slice'`` .
+        ValueError: If `sparse` is False and `target` is ``'CPU'`` .
+        ValueError: If `slice_mode` is ``'field_slice'`` and `feature_num_list` is None.
+        ValueError: If `operator` is not one of ``'SUM'``, ``'MAX'``, ``'MEAN'`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -555,10 +559,11 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
     OPERATOR_MAX = 'MAX'
 
     def __init__(self, vocab_size, embedding_size, field_size, param_init='normal', target='CPU',
-                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM'):
+                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM',
+                 dtype=mstype.float32):
         """Initialize MultiFieldEmbeddingLookup."""
         super(MultiFieldEmbeddingLookup, self).__init__(vocab_size, embedding_size, param_init, target,
-                                                        slice_mode, feature_num_list, max_norm, sparse)
+                                                        slice_mode, feature_num_list, max_norm, sparse, dtype=dtype)
         self.field_size = Validator.check_positive_int(field_size, 'field_size', self.cls_name)
         self.operator = operator
 
@@ -622,8 +627,9 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
                 self.inf_add.shard(((1, 1, get_group_size()), (1, 1, 1)))
         else:
             if is_auto_parallel:
-                raise ValueError("For '{}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' and \
-                                       'table_column_slice'], but got {}".format(self.cls_name, str(slice_mode)))
+                raise ValueError(
+                    f"For '{self.cls_name}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' "
+                    f"and 'table_column_slice'], but got {str(slice_mode)}.")
 
         # Min value for fp32
         self.negative_inf_value = -3.402823466E+38

mindspore/nn/layer/flash_attention.py

@@ -17,12 +17,13 @@ A FlashAttention Layer.
 """
 import math
 
-import mindspore.numpy as mnp
-from mindspore import ops
-from mindspore.common import dtype as mstype
+import mindspore.common.dtype as mstype
 from mindspore.common.tensor import Tensor
+from mindspore import ops
 from mindspore.nn.cell import Cell
 from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_impl import get_flash_attention
+from mindspore.ops.operations.nn_ops import FlashAttentionScore
+from mindspore._c_expression import MSContext
 
 __all__ = ['FlashAttention']
 
@@ -92,6 +93,7 @@ class FlashAttention(Cell):
 
     def __init__(self,
                  head_dim,
+                 head_num,
                  dropout_rate=0.0,
                  prev_block_num=65536,
                  next_block_num=65536,
@@ -104,18 +106,42 @@ class FlashAttention(Cell):
                  ):
         super(FlashAttention, self).__init__()
 
-        self.flash_attention = get_flash_attention(
-            prev_block_num=prev_block_num,
-            next_block_num=next_block_num,
-            tiling_stgy_name=tiling_stgy_name,
-            high_precision=high_precision
-        )
-        self.flash_attention.add_prim_attr("primitive_target", "Ascend")
         scaling_constant = math.sqrt(head_dim)
-        if scaling_constant != 0:
-            self.scale_factor = Tensor([1. / scaling_constant], dtype=mstype.float16)
-        else:
+        if scaling_constant == 0:
             raise ValueError("the scaling constant must not be 0.")
+        self.scale_factor = Tensor([1. / scaling_constant], dtype=mstype.float16)
+
+        self.is_910A = MSContext.get_instance().get_ascend_soc_version() == "Ascend910"
+        if self.is_910A:
+            self.flash_attention = get_flash_attention(
+                prev_block_num=prev_block_num,
+                next_block_num=next_block_num,
+                tiling_stgy_name=tiling_stgy_name,
+                high_precision=high_precision
+            )
+            self.flash_attention.add_prim_attr("primitive_target", "Ascend")
+        else:
+            if alibi:
+                raise ValueError(f"When soc_version is not Ascend910A, alibi must be False")
+            self.transpose_4d_pre = ops.Transpose().shard(((dp, mp, 1, 1),))
+            self.transpose_4d_post = ops.Transpose().shard(((dp, 1, mp, 1),))
+            self.reshape = ops.Reshape()
+            self.zeros_like = ops.ZerosLike().shard(((dp, mp, 1, 1),))
+            self.zeros = ops.Zeros()
+            self.attn_expand_dims = ops.ExpandDims().shard(((dp, 1, 1),))
+            fa_strategies = ((dp, 1, mp),
+                             (dp, 1, mp),
+                             (dp, 1, mp),
+                             (dp, 1, 1, 1))
+            if dropout_rate > 1e-5:
+                fa_strategies += ((dp, mp, 1, 1),)
+            self.flash_attention = FlashAttentionScore(head_num=head_num, pre_tokens=prev_block_num,
+                                                       next_tokens=next_block_num,
+                                                       keep_prob=1 - dropout_rate,
+                                                       scale_value=1.0,
+                                                       inner_precise=0 if high_precision else 1).shard(fa_strategies)
+
+        self.ones = ops.Ones()
         self.dim_mask = Tensor([1 for _ in range(head_dim)], dtype=mstype.int8)
         self.scale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
         self.dropout_rate = dropout_rate
@@ -136,38 +162,35 @@ class FlashAttention(Cell):
                                   such as MatMul. Default: None.
         :return:
         """
-        if in_strategy is not None:
-            shard_stgy = list(in_strategy)
-            shard_stgy.insert(3, (1,))  # dim_mask
-            shard_stgy = tuple(shard_stgy)
-        else:
+        if in_strategy is None:
             # default: dp=1, mp=1, construct inputs only contain query, key, value
-            shard_stgy = (
+            in_strategy = (
                 (1, 1, 1, 1),
                 (1, 1, 1, 1),
                 (1, 1, 1, 1),
-                (1,),  # dim_mask
             )
-        self.flash_attention.shard(shard_stgy)
-        dp = shard_stgy[0][0]
-        mp = shard_stgy[0][1]
+        self.flash_attention.shard(in_strategy)
+        dp = in_strategy[0][0]
+        mp = in_strategy[0][1]
         self.flash_attention.add_prim_attr("dev_matrix_shape", [dp, mp, 1, 1])
         inputs_tensor_map = [
             [3, 2, 1, 0],
             [3, 2, 1, 0],
             [3, 2, 1, 0],
-            [-1]
         ]
         if self.have_attention_mask_batch:
             inputs_tensor_map.append([3, 1, 0])
         else:
             inputs_tensor_map.append([-1, 1, 0])
 
+        input_empty_args_num = 2
         # dropout_mask
         if self.dropout_rate > 1e-5:
+            input_empty_args_num -= 1
             inputs_tensor_map.append([3, 2, 1, 0])
 
         if self.alibi:
+            input_empty_args_num -= 1
             inputs_tensor_map.append([3, 2, 1, 0])
 
         self.flash_attention.add_prim_attr("inputs_tensor_map", inputs_tensor_map)
@@ -178,7 +201,7 @@ class FlashAttention(Cell):
             [3, 2, 1]  # M
         ])
         self.flash_attention.add_prim_attr("as_loss_divisor", 0)
-        self.flash_attention.add_prim_attr("empty_mirror_ops", 1)
+        self.flash_attention.add_prim_attr("empty_mirror_ops", input_empty_args_num)
 
     def construct(self, query, key, value, attn_mask=None, alibi_mask=None):
         """FlashAttention forward
@@ -200,24 +223,42 @@ class FlashAttention(Cell):
         if seq_len % 16 != 0 or k_seq_len % 16 != 0 or k_seq_len != v_seq_len:
             raise ValueError(
                 "query, key, value seq_len must be a multiple of 16, and key seq_len, value seq_len must be the same.")
-        if self.dropout_rate > 1e-5:
-            drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
-            tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
-            ones = self.fill_v2(tensor_shape, self.tensor_one)
-            ones = self.depend(ones, query)
-            drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
-        else:
-            drop_mask = None
+
         if head_dim > 304:
             raise ValueError(
                 "the head_dim must be less than 304, otherwise the ub would be OOM.")
-        if head_dim % 16 != 0:
-            padding_size = 16 - head_dim % 16
-            query = mnp.pad(query, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            key = mnp.pad(key, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            value = mnp.pad(value, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            output, _, _ = self.flash_attention(query, key, value, self.dim_mask, attn_mask, drop_mask, alibi_mask)
-            output = ops.slice(output, [0, 0, 0, 0], [bsz, head_num, seq_len, head_dim])
+
+        if self.is_910A:
+            # 910A -- FlashAttentionPrimtive
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
+                tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
+                ones = self.fill_v2(tensor_shape, self.tensor_one)
+                ones = self.depend(ones, query)
+                drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
+            else:
+                drop_mask = None
+            output, _, _ = self.flash_attention(query, key, value, attn_mask, drop_mask, alibi_mask)
         else:
-            output, _, _ = self.flash_attention(query, key, value, self.dim_mask, attn_mask, drop_mask, alibi_mask)
+            # FlashAttentionScore
+            # Useless input, just for binary calls.
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.reshape(self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob),
+                                              (bsz, head_num, seq_len, seq_len // 8))
+            else:
+                drop_mask_bits = None
+            # (B, N, S, D) -> (B, S, H)
+            query = self.reshape(self.transpose_4d_pre(query, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            key = self.reshape(self.transpose_4d_pre(key, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            value = self.reshape(self.transpose_4d_pre(value, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            attn_mask = self.attn_expand_dims(attn_mask, 1)
+            output, _, _ = self.flash_attention(query,
+                                                key,
+                                                value,
+                                                attn_mask,
+                                                drop_mask_bits,
+                                                None,
+                                                None)
+            output = self.transpose_4d_post(self.reshape(output, (bsz, seq_len, head_num, head_dim)), (0, 2, 1, 3))
+
         return output

mindspore/nn/layer/image.py

@@ -83,17 +83,17 @@ class ImageGradients(Cell):
         _check_input_4d(F.shape(images), "images", self.cls_name)
         batch_size, depth, height, width = P.Shape()(images)
         if height == 1:
-            dy = P.Fill()(P.DType()(images), (batch_size, depth, 1, width), 0)
+            dy = F.fill(P.DType()(images), (batch_size, depth, 1, width), 0)
         else:
             dy = images[:, :, 1:, :] - images[:, :, :height - 1, :]
-            dy_last = P.Fill()(P.DType()(images), (batch_size, depth, 1, width), 0)
+            dy_last = F.fill(P.DType()(images), (batch_size, depth, 1, width), 0)
             dy = P.Concat(2)((dy, dy_last))
 
         if width == 1:
-            dx = P.Fill()(P.DType()(images), (batch_size, depth, height, 1), 0)
+            dx = F.fill(P.DType()(images), (batch_size, depth, height, 1), 0)
         else:
             dx = images[:, :, :, 1:] - images[:, :, :, :width - 1]
-            dx_last = P.Fill()(P.DType()(images), (batch_size, depth, height, 1), 0)
+            dx_last = F.fill(P.DType()(images), (batch_size, depth, height, 1), 0)
             dx = P.Concat(3)((dx, dx_last))
         return dy, dx
 
@@ -571,7 +571,8 @@ class PixelShuffle(Cell):
     <https://arxiv.org/abs/1609.05158>`_ .
 
     Typically, the input is of shape :math:`(*, C \times r^2, H, W)` , and the output is of shape
-    :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor and * is zero or more batch dimensions.
+    :math:`(*, C, H \times r, W \times r)`,
+    where :math:`r` is an upscale factor and :math:`*` is zero or more batch dimensions.
 
     Note:
         The dimension of input Tensor on Ascend should be less than 7.
@@ -621,7 +622,8 @@ class PixelUnshuffle(Cell):
     <https://arxiv.org/abs/1609.05158>`_ .
 
     Typically, the input is of shape :math:`(*, C, H \times r, W \times r)` , and the output is of shape
-    :math:`(*, C \times r^2, H, W)` , where r is a downscale factor and * is zero or more batch dimensions.
+    :math:`(*, C \times r^2, H, W)` ,
+    where :math:`r` is a downscale factor and :math:`*` is zero or more batch dimensions.
 
     Args:
         downscale_factor (int): factor to unshuffle the input, and is a positive integer.

mindspore/nn/layer/math.py

@@ -223,7 +223,6 @@ class LGamma(Cell):
         self.abs = P.Abs()
         self.shape = P.Shape()
         self.dtype = P.DType()
-        self.fill = P.Fill()
         self.floor = P.Floor()
         self.equal = P.Equal()
         self.greater = P.Greater()
@@ -240,7 +239,7 @@ class LGamma(Cell):
         if F.is_sequence_value_unknown(self.shape(x)):
             infinity = self.ones_like(x) * F.cast(self.inf, input_dtype)
         else:
-            infinity = self.fill(input_dtype, self.shape(x), self.inf)
+            infinity = F.fill(input_dtype, self.shape(x), self.inf)
 
         need_to_reflect = self.less(x, 0.5)
         neg_input = -x
@@ -335,7 +334,6 @@ class DiGamma(Cell):
         self.abs = P.Abs()
         self.shape = P.Shape()
         self.dtype = P.DType()
-        self.fill = P.Fill()
         self.floor = P.Floor()
         self.equal = P.Equal()
         self.less = P.Less()
@@ -371,7 +369,7 @@ class DiGamma(Cell):
         reduced_input = x + self.abs(self.floor(x + 0.5))
         reflection = y - self.pi * self.cos(self.pi * reduced_input) / self.sin(self.pi * reduced_input)
         real_result = self.select(need_to_reflect, reflection, y)
-        nan = self.fill(self.dtype(x), self.shape(x), np.nan)
+        nan = F.fill(self.dtype(x), self.shape(x), np.nan)
 
         return self.select(self.logicaland(self.less(x, 0), self.equal(x, self.floor(x))),
                            nan, real_result)
@@ -391,7 +389,6 @@ def _igamma_series(ax, x, a, enabled):
 
     logicaland = P.LogicalAnd()
     greater = P.Greater()
-    fill = P.Fill()
     shape = P.Shape()
     dtype = P.DType()
     select = P.Select()
@@ -424,8 +421,8 @@ def _igamma_series(ax, x, a, enabled):
                 select(enabled, x, vals[4]), select(enabled, dc_da, vals[5]),
                 select(enabled, dans_da, vals[6]))
 
-    ones = fill(dtype(a), shape(a), 1)
-    zeros = fill(dtype(a), shape(a), 0)
+    ones = F.fill(dtype(a), shape(a), 1)
+    zeros = F.fill(dtype(a), shape(a), 0)
     vals = (enabled, a, ones, ones, x, zeros, zeros)
 
     vals = _while_helper_func(cond, body, vals)
@@ -441,7 +438,6 @@ def _igammac_continued_fraction(ax, x, a, enabled):
     greater = P.Greater()
     less = P.Less()
     notequal = P.NotEqual()
-    fill = P.Fill()
     shape = P.Shape()
     dtype = P.DType()
     select = P.Select()
@@ -482,7 +478,7 @@ def _igammac_continued_fraction(ax, x, a, enabled):
         qk_is_nonzero = notequal(qk, 0)
         r = pk / qk
 
-        t = select(qk_is_nonzero, abs_x((ans - r) / r), fill(dtype(t), shape(t), 1))
+        t = select(qk_is_nonzero, abs_x((ans - r) / r), F.fill(dtype(t), shape(t), 1))
         ans = select(qk_is_nonzero, r, ans)
 
         dpk_da = dpkm1_da * z - pkm1 - dpkm2_da * yc + pkm2 * c
@@ -490,7 +486,7 @@ def _igammac_continued_fraction(ax, x, a, enabled):
         dans_da_new = select(qk_is_nonzero, (dpk_da - ans * dqk_da) / qk, dans_da)
         grad_conditional = select(qk_is_nonzero,
                                   abs_x(dans_da_new - dans_da),
-                                  fill(dtype(dans_da), shape(dans_da), 1))
+                                  F.fill(dtype(dans_da), shape(dans_da), 1))
 
         pkm2 = pkm1
         pkm1 = pk
@@ -525,16 +521,16 @@ def _igammac_continued_fraction(ax, x, a, enabled):
 
     y = 1 - a
     z = x + y + 1
-    c = fill(dtype(x), shape(x), 0)
-    pkm2 = fill(dtype(x), shape(x), 1)
+    c = F.fill(dtype(x), shape(x), 0)
+    pkm2 = F.fill(dtype(x), shape(x), 1)
     qkm2 = x
     pkm1 = x + 1
     qkm1 = z * x
     ans = pkm1 / qkm1
-    t = fill(dtype(x), shape(x), 1)
-    dpkm2_da = fill(dtype(x), shape(x), 0)
-    dqkm2_da = fill(dtype(x), shape(x), 0)
-    dpkm1_da = fill(dtype(x), shape(x), 0)
+    t = F.fill(dtype(x), shape(x), 1)
+    dpkm2_da = F.fill(dtype(x), shape(x), 0)
+    dqkm2_da = F.fill(dtype(x), shape(x), 0)
+    dpkm1_da = F.fill(dtype(x), shape(x), 0)
     dqkm1_da = -x
     dans_da = (dpkm1_da - ans * dqkm1_da) / qkm1
     vals = (enabled, ans, t, y, z, c, pkm1, qkm1, pkm2, qkm2, dpkm2_da, dqkm2_da, dpkm1_da, dqkm1_da, dans_da)
@@ -606,7 +602,6 @@ class IGamma(Cell):
         self.exp = P.Exp()
         self.select = P.Select()
         self.zeroslike = P.ZerosLike()
-        self.fill = P.Fill()
         self.shape = P.Shape()
         self.dtype = P.DType()
         self.lgamma = LGamma()
@@ -633,7 +628,7 @@ class IGamma(Cell):
                              1 - _igammac_continued_fraction(ax, x, a, self.logicaland(enabled, use_igammac)),
                              _igamma_series(ax, x, a, self.logicaland(enabled, self.logicalnot(use_igammac))))
         output = self.select(x_is_zero, self.zeroslike(output), output)
-        output = self.select(domain_error, self.fill(self.dtype(a), self.shape(a), np.nan), output)
+        output = self.select(domain_error, F.fill(self.dtype(a), self.shape(a), np.nan), output)
         return output