mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/_grad_experimental/grad_nn_ops.py

@@ -15,7 +15,10 @@
 
 
 """Define the grad rules of neural network related operations."""
+from __future__ import absolute_import
+
 from mindspore import Tensor
+import mindspore.numpy as mnp
 from mindspore.ops import matmul
 from mindspore.ops.operations.nn_ops import GridSampler2D
 from mindspore.ops.operations.nn_ops import GridSampler3D
@@ -35,6 +38,8 @@ from mindspore.ops.operations.nn_ops import FractionalMaxPool3DWithFixedKsize
 from mindspore.ops.operations._grad_ops import FractionalMaxPool3DGradWithFixedKsize
 from mindspore.ops.operations.nn_ops import FractionalAvgPool
 from mindspore.ops.operations._grad_ops import FractionalAvgPoolGrad
+from mindspore.ops.operations.nn_ops import InstanceNormV2
+from mindspore.ops.operations._grad_ops import InstanceNormV2Grad
 from mindspore.ops.operations.nn_ops import MultiMarginLoss
 from mindspore.ops.operations.nn_ops import MultilabelMarginLoss
 from mindspore.ops.operations.nn_ops import NthElement
@@ -56,9 +61,12 @@ from mindspore.ops.operations.nn_ops import ReLUV3
 from mindspore.ops.operations._grad_ops import ReluGrad
 from mindspore.ops.operations.image_ops import ResizeLinear1D
 from mindspore.ops.operations.nn_ops import MaxPool3DWithArgmax
+from mindspore.ops.operations.nn_ops import MaxPoolWithArgmaxV2
 from mindspore.ops.operations.nn_ops import FractionalMaxPoolWithFixedKsize
 from mindspore.ops.operations._grad_ops import FractionalMaxPoolGradWithFixedKsize
 from mindspore.ops.operations.nn_ops import AdaptiveAvgPool3D
+from mindspore.ops.operations.nn_ops import GLU
+from mindspore.ops.operations.nn_ops import AdaptiveMaxPool3D
 
 
 @bprop_getters.register(P.CTCLossV2)
@@ -73,6 +81,24 @@ def get_bprop_ctc_loss_v2(self):
     return bprop
 
 
+@bprop_getters.register(InstanceNormV2)
+def get_bprop_instance_norm_v2(self):
+    """Grad definition for `InstanceNormV2` operation."""
+    grad_ops = InstanceNormV2Grad(self.is_training, self.epsilon)
+
+    def bprop(x, gamma, beta, mean, variance, out, dout):
+        saved_mean = out[1]
+        saved_variance = out[2]
+        grad_ops_out = grad_ops(dout[0], x, gamma, mean, variance, saved_mean, saved_variance)
+        dx = grad_ops_out[0]
+        dgamma = grad_ops_out[1]
+        dbeta = grad_ops_out[2]
+        pd_res = (dx, dgamma, dbeta, zeros_like(mean), zeros_like(variance))
+        return pd_res
+
+    return bprop
+
+
 @bprop_getters.register(P.SoftMarginLoss)
 def get_bprop_soft_margin_loss(self):
     """Grad definition for `SoftMarginLoss` operation."""
@@ -121,7 +147,10 @@ def get_bprop_pad_v3(self):
 
     def bprop(x, paddings, constant_values, out, dout):
         if mode == 'constant':
-            neg_paddings = tuple(-x for x in paddings)
+            if isinstance(paddings, (list, tuple)):
+                neg_paddings = tuple(-x for x in paddings)
+            else:
+                neg_paddings = -paddings
             dx = pad_v3_grad(dout, neg_paddings, zeros_like(constant_values))
         else:
             dx = pad_v3_grad(dout, paddings)
@@ -164,12 +193,10 @@ def get_bprop_celu(self):
     """Grad definition for `CeLU` operation."""
     alpha = self.alpha
     greater_equal = P.GreaterEqual()
-    less = P.Less()
 
     def bprop(x, out, dout):
-        greater = greater_equal(x, 0.0)
-        lesser = less(x, 0.0)
-        dx = dout * (greater * 1.0 + lesser * (out / alpha + 1.0))
+        condition = greater_equal(x, 0.0)
+        dx = dout * mnp.where(condition, 1.0, out / alpha + 1.0)
         return (dx,)
 
     return bprop
@@ -299,10 +326,10 @@ def get_bprop_nth_element(self):
     divide = P.Div()
 
     def bprop(input_x, n, out, dout):
-        indicators = cast(equal(expand_dims(out, -1), input_x), input_x.dtype)
+        indicators = cast(equal(expand_dims(out, -1), input_x), mstype.float32)
         dout = expand_dims(dout, -1)
         num_select = expand_dims(reduce_sum(indicators, -1), -1)
-        return divide(indicators, num_select) * dout, None
+        return cast(divide(indicators, num_select) * dout, input_x.dtype), None
 
     return bprop
 
@@ -354,9 +381,10 @@ def _create_tensor(x_shape):
 def get_bprop_fractional_avg_pool(self):
     """Grad definition for `FractionalAvgPool` operation."""
     fractional_avg_pool_grad = FractionalAvgPoolGrad(overlapping=self.overlapping)
+    shape_op = P.Shape()
 
     def bprop(x, out, dout):
-        dx = fractional_avg_pool_grad(_create_tensor(x.shape), dout[0], out[1], out[2])
+        dx = fractional_avg_pool_grad(_create_tensor(shape_op(x)), dout[0], out[1], out[2])
         return (dx,)
 
     return bprop
@@ -378,6 +406,18 @@ def get_bprop_p_s_r_o_i_pooling(self):
     return bprop
 
 
+@bprop_getters.register(AdaptiveMaxPool3D)
+def get_bprop_adaptive_max_pool_3d(self):
+    """Grad definition for `AdaptiveMaxPool3D` operation."""
+    grad = G.AdaptiveMaxPool3DGrad()
+
+    def bprop(x, output_size, out, dout):
+        dx = grad(dout[0], x, out[1])
+        return (dx, P.ZerosLike()(output_size))
+
+    return bprop
+
+
 @bprop_getters.register(AvgPoolV1)
 def get_bprop_avg_pool_v1_grad(self):
     """Grad definition for `AvgPoolV1` operation."""
@@ -457,6 +497,18 @@ def get_bprop_resize_bilinear(self):
     return bprop
 
 
+@bprop_getters.register(GLU)
+def get_bprop_glu(self):
+    """Grad definition for `Glu` operation."""
+    input_grad = G.GluGrad(self.axis)
+
+    def bprop(x, out, dout):
+        dx = input_grad(dout, x)
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(MaxPool3DWithArgmax)
 def get_bprop_maxpool3dwithargmax(self):
     """Grad definition for `MaxPool3DWithArgmax` operation."""
@@ -475,6 +527,24 @@ def get_bprop_maxpool3dwithargmax(self):
     return bprop
 
 
+@bprop_getters.register(MaxPoolWithArgmaxV2)
+def get_bprop_maxpoolwithargmaxv2(self):
+    """Grad definition for `MaxPoolWithArgmaxV2` operation."""
+    maxpoolwithargmaxv2_grad = G.MaxPoolGradWithArgmaxV2(
+        kernel_size=self.kernel_size,
+        strides=self.strides,
+        pads=self.pads,
+        dilation=self.dilation,
+        ceil_mode=self.ceil_mode,
+        argmax_type=self.argmax_type)
+
+    def bprop(x, out, dout):
+        dx = maxpoolwithargmaxv2_grad(x, dout[0], out[1])
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(FractionalMaxPoolWithFixedKsize)
 def get_bprop_fractional_max_pool_with_fixed_ksize(self):
     """Grad definition for 'FractionalMaxPoolWithFixedKsize' operation."""

mindspore/ops/_grad_experimental/grad_scalar_ops.py

@@ -0,0 +1,112 @@
+# Copyright 2023 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.
+# ============================================================================
+
+"""Generate bprop for quantization aware ops"""
+
+from mindspore.ops.operations import _scalar_ops
+from mindspore.ops._grad.grad_base import bprop_getters
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+
+
+@bprop_getters.register(_scalar_ops.ScalarAdd)
+def get_bprop_scalar_add(self):
+    """Grad definition for `ScalarAdd` operation."""
+
+    def bprop(x, y, out, dout):
+        return dout, dout
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarSub)
+def get_bprop_scalar_sub(self):
+    """Grad definition for `ScalarSub` operation."""
+
+    def bprop(x, y, out, dout):
+        return dout, 0 - dout
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarMul)
+def get_bprop_scalar_mul(self):
+    """Grad definition for `ScalarMul` operation."""
+
+    def bprop(x, y, out, dout):
+        bc_dx = y * dout
+        bc_dy = x * dout
+        return bc_dx, bc_dy
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarDiv)
+def get_bprop_scalar_div(self):
+    """Grad definition for `ScalarDiv` operation."""
+
+    def bprop(x, y, out, dout):
+        bc_dx = dout / y
+        bc_dy = 0 - bc_dx * out
+        return bc_dx, bc_dy
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarFloordiv)
+def get_bprop_scalar_floordiv(self):
+    """Grad definition for `ScalarFloorDiv` operation."""
+
+    def bprop(x, y, out, dout):
+        return zeros_like(x), zeros_like(y)
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarMod)
+def get_bprop_scalar_mod(self):
+    """Grad definition for `ScalarMod` operation."""
+
+    def bprop(x, y, out, dout):
+        bc_dx = dout
+        bc_dy = -dout * (x // y)
+        return bc_dx, bc_dy
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.scalar_eq)
+@bprop_getters.register(_scalar_ops.scalar_le)
+@bprop_getters.register(_scalar_ops.scalar_lt)
+@bprop_getters.register(_scalar_ops.scalar_ge)
+@bprop_getters.register(_scalar_ops.scalar_gt)
+@bprop_getters.register(_scalar_ops.bit_and)
+@bprop_getters.register(_scalar_ops.bit_or)
+def get_bprop_scalar_logic(self):
+    """Grad definition for `ScalarLogicOps` operation."""
+
+    def bprop(x, y, out, dout):
+        return zeros_like(x), zeros_like(y)
+
+    return bprop
+
+
+@bprop_getters.register(_scalar_ops.ScalarBool)
+def get_bprop_scalar_bool(self):
+    """Grad definition for `ScalarBool` operation."""
+
+    def bprop(x, out, dout):
+        return zeros_like(x)
+
+    return bprop

mindspore/ops/_grad_experimental/grad_sparse_ops.py

@@ -14,6 +14,7 @@
 # ============================================================================
 
 """bprop primitives"""
+from mindspore.ops.operations.sparse_ops import RaggedTensorToSparse
 from mindspore.ops.operations.sparse_ops import CSRSparseMatrixToSparseTensor
 from mindspore.ops.operations.sparse_ops import SparseReorder
 from mindspore.ops.operations.sparse_ops import SparseTensorToCSRSparseMatrix
@@ -24,7 +25,16 @@ from mindspore.ops.operations.sparse_ops import SparseSegmentSum
 from mindspore.ops.operations.sparse_ops import SparseSegmentSumWithNumSegments
 from mindspore.ops.operations.sparse_ops import SparseSegmentSqrtN
 from mindspore.ops.operations.sparse_ops import SparseSegmentSqrtNWithNumSegments
+from mindspore.ops.operations.sparse_ops import SparseFillEmptyRows
 from mindspore.ops.operations.sparse_ops import SparseSegmentMeanWithNumSegments
+from mindspore.ops.operations.sparse_ops import SparseSlice
+from mindspore.ops.operations.sparse_ops import SparseDenseCwiseMul
+from mindspore.ops.operations.sparse_ops import SparseDenseCwiseDiv
+from mindspore.ops.operations.sparse_ops import SparseTensorDenseAdd
+from mindspore.ops.operations.sparse_ops import Sspaddmm
+from mindspore.ops.operations._inner_ops import IsSubClass
+import mindspore as ms
+from mindspore.ops.operations import _map_tensor_ops
 from mindspore.common import dtype as mstype
 from mindspore import Tensor
 from mindspore.ops.primitive import constexpr
@@ -33,10 +43,77 @@ from mindspore.ops import operations as P
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops._grad.grad_base import bprop_getters
-from mindspore.ops._utils.utils import is_shape_unknown
+from mindspore import context
 
 # Unused parameters are placeholders.
 dyn_shape_op = P.TensorShape()
+is_sub_class = IsSubClass()
+
+
+@bprop_getters.register(SparseDenseCwiseMul)
+def get_bprop_sparse_dense_cwise_mul(self):
+    """Grad definition for 'SparseDenseCwiseMul' operation"""
+
+    slice_op = P.Slice()
+    ones = P.Ones()
+    shape = P.Shape()
+    cast = P.Cast()
+    gather_nd = P.GatherNd()
+    size_op = P.Size()
+    to_array = P.TupleToArray()
+    sparse_tensor_dense_add = SparseTensorDenseAdd()
+
+    def bprop(x1_indices, x1_values, x1_shape, x2, out, dout):
+        x2_shape = cast(to_array(shape(x2)), mstype.int64)
+        augmented_x2_shape = P.Concat(0)((ones(size_op(x1_shape) - size_op(x2_shape), mstype.int64), x2_shape))
+        scaling = x1_shape // augmented_x2_shape
+        scaled_indices = x1_indices // scaling
+        scaled_indices = cast(slice_op(scaled_indices, [0, size_op(x1_shape) - size_op(x2_shape)], [-1, -1]),
+                              mstype.int64)
+        dense_vals = gather_nd(x2, scaled_indices)
+        dx1 = dout * dense_vals
+        dx2_val = dout * x1_values
+        dx2 = sparse_tensor_dense_add(scaled_indices, dx2_val, x2_shape, zeros_like(x2))
+        d_all = (zeros_like(x1_indices), dx1, zeros_like(x1_shape), dx2)
+        return d_all
+
+    return bprop
+
+
+@bprop_getters.register(SparseDenseCwiseDiv)
+def get_bprop_sparse_dense_cwise_div(self):
+    """Grad definition for 'SparseDenseCwiseDiv' operation"""
+
+    slice_op = P.Slice()
+    ones = P.Ones()
+    shape = P.Shape()
+    cast = P.Cast()
+    gather_nd = P.GatherNd()
+    size_op = P.Size()
+    to_array = P.TupleToArray()
+    sparse_tensor_dense_add = SparseTensorDenseAdd()
+
+    def bprop(x1_indices, x1_values, x1_shape, x2, out, dout):
+
+        x2_shape = cast(to_array(shape(x2)), mstype.int64)
+        augmented_x2_shape = P.Concat(0)((ones(size_op(x1_shape) - size_op(x2_shape), mstype.int64), x2_shape))
+        scaling = x1_shape // augmented_x2_shape
+        scaled_indices = x1_indices // scaling
+        scaled_indices = cast(slice_op(scaled_indices, [0, size_op(x1_shape) - size_op(x2_shape)],
+                                       [-1, -1]), mstype.int64)
+        dense_vals = gather_nd(x2, scaled_indices)
+        dx1 = dout / dense_vals
+        dense_vals_2 = dense_vals * dense_vals
+        w = x1_values / dense_vals_2
+        w = (-1) * w
+        dx2_val = dout * w
+        dx2 = sparse_tensor_dense_add(scaled_indices, dx2_val, x2_shape,
+                                      zeros_like(x2))
+        d_all = (zeros_like(x1_indices), dx1, zeros_like(x1_shape), dx2)
+        return d_all
+
+    return bprop
+
 
 
 @constexpr
@@ -56,7 +133,7 @@ def get_bprop_sparse_softmax(self):
         default_values = _create_tensor(0, values.dtype)
         out_dout = mul(out, dout)
         sp_product = sparse_to_dense(indices, shape, out_dout, default_values)
-        sum_reduced = -reduce_sum(sp_product, -1)
+        sum_reduced = -1 * reduce_sum(sp_product, -1)
         sp_sum = sparse_dense_cwise_add(indices, dout, shape, sum_reduced)
         grad_x = mul(sp_sum, out)
         return zeros_like(indices), grad_x, zeros_like(shape)
@@ -111,7 +188,7 @@ def get_bprop_sparse_segment_sqrt_n(self):
 
     def bprop(x, indices, segment_ids, out, dout):
         shape_x = shape(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             shape_x = dyn_shape_op(x)
         output_dim0 = P.Cast()(shape_x[0], mstype.int32)
         indices = F.cast(indices, mstype.int32)
@@ -131,7 +208,7 @@ def get_bprop_sparse_segment_sqrt_n_with_num_segments(self):
 
     def bprop(x, indices, segment_ids, num_segments, out, dout):
         shape_x = shape(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             shape_x = dyn_shape_op(x)
         output_dim0 = P.Cast()(shape_x[0], mstype.int32)
         indices = F.cast(indices, mstype.int32)
@@ -146,12 +223,14 @@ def get_bprop_sparse_segment_sqrt_n_with_num_segments(self):
 @bprop_getters.register(SparseSegmentSum)
 def get_bprop_sparse_segment_sum(self):
     """Grad definition for `SparseSegmentSum` operation."""
-    input_grad = G.SparseSegmentSumGrad()
+    gather = P.Gather()
+    unsorted_segment_sum = P.UnsortedSegmentSum()
     shape = P.Shape()
 
-    def bprop(x, indices, segment_ids, out, dout):
+    def bprop_gpu(x, indices, segment_ids, out, dout):
+        input_grad = G.SparseSegmentSumGrad()
         shape_x = shape(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             shape_x = dyn_shape_op(x)
         output_dim0 = P.Cast()(shape_x[0], mstype.int32)
         indices = F.cast(indices, mstype.int32)
@@ -160,18 +239,46 @@ def get_bprop_sparse_segment_sum(self):
         all_d = (dx, zeros_like(indices), zeros_like(segment_ids))
         return all_d
 
+    def bprop(x, indices, segment_ids, out, dout):
+        shape_x = dyn_shape_op(x)
+        output_dim0 = P.Cast()(shape_x[0], mstype.int32)
+        segment_ids = F.cast(segment_ids, mstype.int32)
+        input0 = gather(dout, segment_ids, 0)
+        input0 = F.cast(input0, mstype.float32)
+        indices = F.cast(indices, mstype.int32)
+        dx = unsorted_segment_sum(input0, indices, output_dim0)
+        dx = F.cast(dx, F.dtype(dout))
+        return dx, zeros_like(indices), zeros_like(segment_ids)
+
+    if context.get_context('device_target') == "GPU":
+        return bprop_gpu
+
+    return bprop
+
+
+@bprop_getters.register(SparseFillEmptyRows)
+def get_bprop_sparsefillemptyrows(self):
+    """Grad definition for `SparseFillEmptyRows` operation."""
+    op = G.SparseFillEmptyRowsGrad()
+
+    def bprop(indices, values, dense_shape, default_value, out, dout):
+        dx = op(out[3], dout[1])
+        dx_all = (zeros_like(indices), dx[0], zeros_like(dense_shape), dx[1])
+        return dx_all
     return bprop
 
 
 @bprop_getters.register(SparseSegmentSumWithNumSegments)
 def get_bprop_sparse_segment_sum_with_num_segments(self):
     """Grad definition for `SparseSegmentSumWithNumSegments` operation."""
-    input_grad = G.SparseSegmentSumGrad()
+    gather = P.Gather()
+    unsorted_segment_sum = P.UnsortedSegmentSum()
     shape = P.Shape()
 
-    def bprop(x, indices, segment_ids, num_segments, out, dout):
+    def bprop_gpu(x, indices, segment_ids, num_segments, out, dout):
+        input_grad = G.SparseSegmentSumGrad()
         shape_x = shape(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             shape_x = dyn_shape_op(x)
         output_dim0 = P.Cast()(shape_x[0], mstype.int32)
         indices = F.cast(indices, mstype.int32)
@@ -180,6 +287,21 @@ def get_bprop_sparse_segment_sum_with_num_segments(self):
         all_d = (dx, zeros_like(indices), zeros_like(segment_ids), zeros_like(num_segments))
         return all_d
 
+    def bprop(x, indices, segment_ids, num_segments, out, dout):
+        shape_x = dyn_shape_op(x)
+        output_dim0 = P.Cast()(shape_x[0], mstype.int32)
+        segment_ids = F.cast(segment_ids, mstype.int32)
+        input0 = gather(dout, segment_ids, 0)
+        input0 = F.cast(input0, mstype.float32)
+        indices = F.cast(indices, mstype.int32)
+        dx = unsorted_segment_sum(input0, indices, output_dim0)
+        dx = F.cast(dx, F.dtype(dout))
+        all_d = (dx, zeros_like(indices), zeros_like(segment_ids), zeros_like(num_segments))
+        return all_d
+
+    if context.get_context('device_target') == "GPU":
+        return bprop_gpu
+
     return bprop
 
 
@@ -191,7 +313,7 @@ def get_bprop_sparse_segment_mean_with_num_segments(self):
 
     def bprop(x, indices, segment_ids, num_segments, out, dout):
         x_shp = shape(x)
-        if is_shape_unknown(x_shp):
+        if F.is_sequence_value_unknown(x_shp):
             x_shp = dyn_shape_op(x)
             output_dim0 = F.cast(x_shp[0], mstype.int32)
         else:
@@ -205,6 +327,56 @@ def get_bprop_sparse_segment_mean_with_num_segments(self):
     return bprop
 
 
+@bprop_getters.register(SparseSlice)
+def get_bprop_sparse_slice(self):
+    """Grad definition for `SparseSlice` operation."""
+    sparse_slice_grad = G.SparseSliceGrad()
+
+    def bprop(indices, values, shape, start, size, out, dout):
+        grad_op = sparse_slice_grad(dout[1], indices, start, out[0])
+        result_all = (zeros_like(indices), grad_op, zeros_like(shape), zeros_like(start), zeros_like(size))
+        return result_all
+
+    return bprop
+
+
+@bprop_getters.register(Sspaddmm)
+def get_bprop_sspaddmm(self):
+    """Grad definition for `Sspaddmm` operation."""
+    def bprop(x1_indices, x1_values, x1_shape, x2_indices, x2_values, x2_shape, x3_dense, alpha, beta, out, dout):
+        dx_all = (zeros_like(x1_indices), zeros_like(x1_values), zeros_like(x1_shape),
+                  zeros_like(x2_indices), zeros_like(x2_values), zeros_like(x2_shape),
+                  zeros_like(x3_dense), zeros_like(alpha), zeros_like(beta))
+        return dx_all
+
+    return bprop
+
+
+@bprop_getters.register(RaggedTensorToSparse)
+def get_bprop_ragged_tensor_to_sparse(self):
+    """Grad definition for `RaggedTensorToSparse` operation."""
+    shape = P.Shape()
+    reshape = P.Reshape()
+
+    def bprop(rt_nested_splits, rt_dense_values, out, dout):
+        ragged_values_shape = shape(rt_dense_values)
+        ragged_values_grad = reshape(dout[1], ragged_values_shape)
+
+        if is_sub_class(F.typeof(rt_nested_splits), ms.list_):
+            split = []
+            for i in enumerate(rt_nested_splits):
+                split.append(zeros_like(i))
+            all_d = (split, ragged_values_grad)
+            return all_d
+        split = ()
+        for i in enumerate(rt_nested_splits):
+            split = split + (zeros_like(i),)
+        all_d = (split, ragged_values_grad)
+        return all_d
+
+    return bprop
+
+
 @bprop_getters.register(SparseReorder)
 def get_bprop_sparse_reorder(self):
     """Grad definition for `SparseReorder` operation."""
@@ -215,11 +387,23 @@ def get_bprop_sparse_reorder(self):
     def bprop(indices, values, shape, out, dout):
         num_entries = F.shape(indices)[0]
         start = Tensor(0, dtype=mstype.int32)
-        limit = Tensor(num_entries, dtype=mstype.int32)
+        limit = P.Cast()(num_entries, mstype.int32)
         delta = Tensor(1, dtype=mstype.int32)
         entry_indices = range_op(start, limit, delta)
         output = sparse_reorder_op(indices, entry_indices, shape)
         inverted_permutation = F.sort(output[1].astype(mstype.float32))[1]
         axis = 0
-        return None, gather_op(dout[1], inverted_permutation, axis), None
+        return zeros_like(indices), gather_op(dout[1], inverted_permutation, axis), zeros_like(shape)
+
+    return bprop
+
+
+@bprop_getters.register(_map_tensor_ops.MapTensorGet)
+def get_bprop_map_tensor_get(self):
+    """Grad definition for `MapTensorGet` operation."""
+    grad_op = G.MapTensorGetGrad()
+
+    def bprop(map_tensor, key_tensor, out, dout):
+        grad_map_tensor = grad_op(map_tensor, key_tensor, dout)
+        return grad_map_tensor, zeros_like(key_tensor)
     return bprop

mindspore/ops/_op_impl/__init__.py

@@ -15,10 +15,10 @@
 """Operators info register."""
 
 import platform
-from .aicpu import *
-from .cpu import *
+from mindspore.ops._op_impl.aicpu import *
+from mindspore.ops._op_impl.cpu import *
 if "Windows" not in platform.system():
+    from mindspore.ops._op_impl.akg import *
     from mindspore.ops._op_impl.tbe import *
-    from .akg import *
 
 __all__ = []

mindspore/ops/_op_impl/_custom_op/__init__.py

@@ -14,7 +14,6 @@
 # ============================================================================
 
 """custom ops"""
-from __future__ import absolute_import
 from mindspore.ops._op_impl._custom_op.dsd_impl import dsd_matmul
 from mindspore.ops._op_impl._custom_op.dsd_back_impl import dsdbpropimpl
 from .batchnorm_fold import _batchnorm_fold_tbe

mindspore/ops/_op_impl/_custom_op/_basic.py

@@ -46,7 +46,6 @@ def _check_mn_shape(m_shape, n_shape, km_shape, kn_shape):
         raise RuntimeError("input shape M and N can't both be 1")
 
     if km_shape != kn_shape:
-        print(km_shape, kn_shape)
         raise RuntimeError("reduce axis not same")
 
     if m_shape % cce.BLOCK_IN != 0 and m_shape != 1:

mindspore/ops/_op_impl/_custom_op/batch_matmul_impl.py

@@ -16,7 +16,7 @@
 from __future__ import absolute_import
 
 from te import tik
-from topi.cce import util
+from tbe.tvm.topi.cce import util
 
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 

mindspore/ops/_op_impl/_custom_op/batchnorm_fold.py

@@ -18,8 +18,8 @@ from __future__ import absolute_import
 
 import te
 from te import tvm
-from topi import generic
-from topi.cce import util
+from tbe.tvm.topi import generic
+from tbe.tvm.topi.cce import util
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
 batch_norm_op_info = TBERegOp("BatchNormFoldD") \
@@ -62,6 +62,8 @@ def _batchnorm_fold_compute(x_input, x_sum, x_square_sum, mean, variance, moment
     """_batchnorm_fold_compute"""
     shape_x = te.lang.cce.util.shape_to_list(x_input.shape)
     num = shape_x[0] * shape_x[2] * shape_x[3]
+    if num == 0.0:
+        raise ValueError('`num` is zero, which leads to divide zero error.')
     num_rec = 1.0 / num
 
     # compute the mean of x

mindspore/ops/_op_impl/_custom_op/batchnorm_fold2.py

@@ -17,8 +17,8 @@
 import te.lang.cce
 from te import tvm
 from te.platform.fusion_manager import fusion_manager
-from topi import generic
-from topi.cce import util
+from tbe.tvm.topi import generic
+from tbe.tvm.topi.cce import util
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
 SHAPE_SIZE_LIMIT = 2147483648

mindspore/ops/_op_impl/_custom_op/batchnorm_fold2_grad.py

@@ -17,8 +17,8 @@
 import te.lang.cce
 from te import tvm
 from te.platform.fusion_manager import fusion_manager
-from topi import generic
-from topi.cce import util
+from tbe.tvm.topi import generic
+from tbe.tvm.topi.cce import util
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
 SHAPE_SIZE_LIMIT = 2147483648