mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/_grad/grad_sequence_ops.py

@@ -0,0 +1,296 @@
+# 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.
+# ============================================================================
+
+"""grad_sequence_ops"""
+
+from mindspore.ops.operations import _sequence_ops as seq
+from mindspore.ops import operations as P
+from mindspore.ops import functional as F
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+from mindspore.ops._grad.grad_base import bprop_getters
+from mindspore.ops.primitive import Primitive
+
+
+tuple_setitem = Primitive("tuple_setitem")
+list_setitem = Primitive("list_setitem")
+
+
+@bprop_getters.register(seq.SequenceCount)
+def get_bprop_count(self):
+    """Generate bprop for SequenceCount"""
+
+    def bprop(x, y, out, dout):
+        return (zeros_like(x), zeros_like(y))
+
+    return bprop
+
+
+@bprop_getters.register(seq.sequence_len)
+def get_bprop_sequence_len(self):
+    """Generate bprop for sequence_len"""
+    def bprop(x, out, dout):
+        return (zeros_like(x),)
+
+    return bprop
+
+
+@bprop_getters.register(seq.SequenceAdd)
+def get_bprop_sequence_add(self):
+    """Generate bprop for SequenceAdd"""
+    def bprop(x, y, out, dout):
+        out_offset = seq.SequenceAddOffset()(x, y)
+        dx = seq.SequenceSlice()(dout, out_offset[0], len(x), 1)
+        dy = seq.SequenceSlice()(dout, out_offset[1], len(x) + len(y), 1)
+
+        return (dx, dy)
+
+    return bprop
+
+
+@bprop_getters.register(seq.SequenceSlice)
+def get_bprop_slice(self):
+    """Generate bprop for SequenceSlice"""
+
+    def bprop(x, start, stop, step, out, dout):
+        dx = seq.SequenceSliceGrad()(dout, x, start, stop, step)
+        return (dx, zeros_like(start), zeros_like(stop), zeros_like(step))
+
+    return bprop
+
+
+@bprop_getters.register(seq.SequenceIndex)
+def get_bprop_index(self):
+    """Generate bprop for SequenceIndex"""
+
+    def bprop(x, y, start, end, out, dout):
+        return (zeros_like(x), zeros_like(y), zeros_like(start), zeros_like(end))
+
+    return bprop
+
+
+@bprop_getters.register(seq.InSequence)
+def get_bprop_insequence(self):
+    """Generate bprop for InSequence"""
+
+    def bprop(x, y, out, dout):
+        return (zeros_like(x), seq.SequenceZerosLike()(y))
+
+    return bprop
+
+
+@bprop_getters.register("tuple_equal")
+@bprop_getters.register("list_equal")
+def get_bprop_seq_equal(self):
+    """Generate bprop for tuple_equal and list_equal"""
+
+    def bprop(x, y, out, dout):
+        return (zeros_like(x), zeros_like(y))
+
+    return bprop
+
+
+@bprop_getters.register("shape_mul")
+def get_bprop_shape_mul(self):
+    """Generate bprop for tuple_equal and list_equal"""
+
+    def bprop(x, out, dout):
+        dx = seq.ShapeMulGrad()(x, dout)
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register("tuple_setitem")
+def get_bprop_tuple_setitem(self):
+    """Generate bprop for TupleSetItem and ListSetItem"""
+
+    def bprop(x, idx, value, out, dout):
+        d_x = tuple_setitem(dout, idx, zeros_like(value))
+        d_value = dout[idx]
+        d_idx = 0
+        return (d_x, zeros_like(d_idx), d_value)
+
+    return bprop
+
+
+@bprop_getters.register("list_setitem")
+def get_bprop_lsit_setitem(self):
+    """Generate bprop for TupleSetItem and ListSetItem"""
+
+    def bprop(x, idx, value, out, dout):
+        d_x = list_setitem(dout, idx, zeros_like(value))
+        d_value = dout[idx]
+        d_idx = 0
+        return (d_x, zeros_like(d_idx), d_value)
+
+    return bprop
+
+
+@bprop_getters.register(seq.ListAppend)
+def get_bprop_list_append(self):
+    """Generate bprop for ListAppend"""
+
+    def bprop(x, value, out, dout):
+        d_x = seq.ListAppendAndInsertGrad()(dout, -1)
+        return (d_x, zeros_like(value))
+
+    return bprop
+
+
+@bprop_getters.register(seq.ListInsert)
+def get_bprop_list_insert(self):
+    """Generate bprop for ListInsert"""
+
+    def bprop(x, idx, value, out, dout):
+        d_x = seq.ListAppendAndInsertGrad()(dout, idx)
+        return (d_x, zeros_like(idx), zeros_like(value))
+
+    return bprop
+
+
+@bprop_getters.register(seq.TupleToTensor)
+def get_bprop_tuple_to_tensor(self):
+    """Generate bprop for TupleToTensor"""
+
+    def bprop(x, dtype, out, dout):
+        tuple_type = F.typeof(x)
+        dout = P.Cast()(dout, tuple_type)
+        d_x = seq.TensorToTuple()(dout)
+        return (d_x, zeros_like(dtype))
+
+    return bprop
+
+
+@bprop_getters.register(seq.ListToTensor)
+def get_bprop_list_to_tensor(self):
+    """Generate bprop for ListToTensor"""
+
+    def bprop(x, dtype, out, dout):
+        tuple_type = F.typeof(x)
+        dout = P.Cast()(dout, tuple_type)
+        d_x = seq.TensorToList()(dout)
+        return (d_x, zeros_like(dtype))
+
+    return bprop
+
+
+@bprop_getters.register(P.ScalarToTensor)
+def get_bprop_scalar_to_tensor(self):
+    """Generate bprop for ScalarToTensor"""
+
+    def bprop(x, dtype, out, dout):
+        scalar_type = F.typeof(x)
+        dout = P.Cast()(dout, scalar_type)
+        d_x = seq.TensorToScalar()(dout)
+        return (d_x, zeros_like(dtype))
+
+    return bprop
+
+
+@bprop_getters.register(seq.TensorToTuple)
+def get_bprop_tensor_to_tuple(self):
+    """Generate bprop for TensorToTuple"""
+
+    def bprop(x, out, dout):
+        dtype = F.typeof(x)
+        d_x = seq.TupleToTensor()(dout, dtype)
+        return (d_x,)
+
+    return bprop
+
+
+@bprop_getters.register(seq.TensorToList)
+def get_bprop_tensor_to_list(self):
+    """Generate bprop for TensorToList"""
+
+    def bprop(x, out, dout):
+        dtype = F.typeof(x)
+        d_x = seq.ListToTensor()(dout, dtype)
+        return (d_x,)
+
+    return bprop
+
+
+@bprop_getters.register(seq.TensorToScalar)
+def get_bprop_tensor_to_scalar(self):
+    """Generate bprop for TensorToScalar"""
+
+    def bprop(x, out, dout):
+        dtype = F.typeof(x)
+        d_x = P.ScalarToTensor()(dout, dtype)
+        return (d_x,)
+
+    return bprop
+
+
+@bprop_getters.register("tuple_le")
+@bprop_getters.register("tuple_lt")
+@bprop_getters.register("list_le")
+@bprop_getters.register("list_lt")
+def get_bprop_less(self):
+    """Generate bprop for SequenceLessThan and SequenceLessEqual"""
+
+    def bprop(x, y, out, dout):
+        return (zeros_like(x), zeros_like(y))
+
+    return bprop
+
+
+@bprop_getters.register(seq.SequenceMul)
+def get_bprop_mul(self):
+    """Generate bprop for SequenceMul"""
+
+    def bprop(x, y, out, dout):
+        dx = x
+        if isinstance(x, tuple):
+            for i in range(len(x)):
+                dx = tuple_setitem(dx, i, dout[i])
+        else:
+            for i in range(len(x)):
+                dx = list_setitem(dx, i, dout[i])
+        return (dx, zeros_like(y))
+
+    return bprop
+
+
+@bprop_getters.register(seq.SequenceMin)
+@bprop_getters.register(seq.SequenceMax)
+def get_bprop_max_min(self):
+    """Generate bprop for SequenceMax and SequenceMax"""
+
+    def bprop(x, out, dout):
+        index = x.index(out)
+        if isinstance(x, tuple):
+            dx = tuple_setitem(zeros_like(x), index, dout)
+        else:
+            dx = list_setitem(zeros_like(x), index, dout)
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register("tuple_greater_than")
+@bprop_getters.register("list_greater_than")
+@bprop_getters.register("tuple_greater_equal")
+@bprop_getters.register("list_greater_equal")
+def get_bprop_greater(self):
+    """Generate bprop for tuple_greater_than, list_greater_than,
+    tuple_greater_equal, list_greater_equal.
+    """
+
+    def bprop(x, y, out, dout):
+        return (zeros_like(x), zeros_like(y))
+
+    return bprop

mindspore/ops/_grad/grad_sparse.py

@@ -14,7 +14,6 @@
 # ============================================================================
 
 """bprop primitives"""
-from mindspore.ops._utils.utils import is_shape_unknown
 from mindspore.ops._grad.grad_base import bprops, bprop_getters
 from mindspore.ops.composite.multitype_ops._constexpr_utils import infer_out_shape
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
@@ -122,7 +121,7 @@ def get_bprop_sparse_add(self):
         dx1, dx2 = sparse_add_grad(dout[1], x1_indices, x2_indices, out[0])
         ret0 = zeros_like(x1_indices)
         shp = shape_op(x1_values)
-        if is_shape_unknown(shp):
+        if F.is_sequence_value_unknown(shp):
             shp = dyn_shape_op(x1_values)
         dx1_shape = shp
         ret1 = reshape(dx1, dx1_shape)
@@ -130,7 +129,7 @@ def get_bprop_sparse_add(self):
 
         ret3 = zeros_like(x2_indices)
         shp = shape_op(x2_values)
-        if is_shape_unknown(shp):
+        if F.is_sequence_value_unknown(shp):
             shp = dyn_shape_op(x2_values)
         dx2_shape = shp
         ret4 = reshape(dx2, dx2_shape)
@@ -223,7 +222,7 @@ def get_bprop_csr_mul(self):
     to index the dense input.
     """
     def bprop(indptr, indices, values, shape, dense, out, dout):
-        csr_tensor_grad_value = F.csr_mul(F.make_csr_tensor(indptr, indices, dout, shape), dense)
+        csr_tensor_grad_value = F.csr_mul(F.make_csr_tensor(indptr, indices, dout, shape), dense).values
         dense_grad_value = F.mul(dout, values)
         dense_grad = F.make_csr_tensor(indptr, indices, dense_grad_value, shape)
         if len(dense.shape) == 1 or dense.shape[0] == 1:
@@ -261,9 +260,11 @@ def get_bprop_csr_div(self):
         shape_y = feature_dim + shape[batch_dim_dense_start:]
         reduce_x, reduce_y = F.broadcast_gradient_args(shape_x, shape_y)
 
-        csr_tensor_grad_value = F.csr_div(F.make_csr_tensor(indptr, indices, dout, shape), dense)
+        csr_tensor_grad = F.csr_div(F.make_csr_tensor(indptr, indices, dout, shape), dense)
         if reduce_x:
-            csr_tensor_grad_value = P.ReduceSum(True)(csr_tensor_grad_value, reduce_x)
+            csr_tensor_grad_value = P.ReduceSum(True)(csr_tensor_grad.values, reduce_x)
+        else:
+            csr_tensor_grad_value = csr_tensor_grad.values
         dense_grad_value = F.neg_tensor(F.mul(out, csr_tensor_grad_value))
         dense_grad = F.make_csr_tensor(indptr, indices, dense_grad_value, shape)
         if len(dense.shape) == 1 or dense.shape[0] == 1:

mindspore/ops/_grad_experimental/__init__.py

@@ -24,5 +24,6 @@ from mindspore.ops._grad_experimental import grad_math_ops
 from mindspore.ops._grad_experimental import grad_linalg_ops
 from mindspore.ops._grad_experimental import grad_sparse
 from mindspore.ops._grad_experimental import grad_sparse_ops
+from mindspore.ops._grad_experimental import grad_scalar_ops
 
 __all__ = ['get_bprop_fn']

mindspore/ops/_grad_experimental/grad_array_ops.py

@@ -33,6 +33,7 @@ from mindspore.ops.operations.array_ops import Mvlgamma
 from mindspore.ops.operations.array_ops import Triu
 from mindspore.ops.operations.array_ops import IdentityN
 from mindspore.ops.operations.array_ops import IndexFill
+from mindspore.ops.operations.array_ops import IndexPut
 from mindspore.ops.operations.array_ops import CheckNumerics
 from mindspore.ops.operations.array_ops import ConjugateTranspose
 from mindspore.ops.operations.array_ops import SegmentMax
@@ -47,13 +48,48 @@ from mindspore.ops.operations.array_ops import Im2Col
 from mindspore.ops.operations.array_ops import Col2Im
 from mindspore.ops.operations.array_ops import StridedSliceV2
 from mindspore.ops.operations.array_ops import MaskedScatter
+from mindspore.ops.operations.array_ops import MaskedSelect
+from mindspore.ops.operations.array_ops import CountNonZero
 from mindspore.ops.operations._grad_ops import StridedSliceV2Grad
 from mindspore.ops.operations.random_ops import LogNormalReverse
+from mindspore.ops.operations.random_ops import ParameterizedTruncatedNormal
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
-from mindspore.ops._utils.utils import is_shape_unknown
 from mindspore.ops.operations import _grad_ops as G
+from mindspore import context
+
+
+@constexpr
+def _raise_value_error(*info):
+    info_str = ""
+    for obj in info:
+        info_str = info_str + f"{obj}"
+    raise ValueError(info_str)
+
+
+@bprop_getters.register(P.FillV2)
+def get_bprop_fill_v2(self):
+    """Generate bprop for FillV2"""
+    sum_op = P.ReduceSum()
+    cast_op = P.Cast()
+    shape_op = P.TensorShape()
+
+    def bprop(shape, value, out, dout):
+        dout_type = F.dtype(dout)
+        type_list = [
+            mstype.int8, mstype.int16, mstype.int32, mstype.int64,
+            mstype.uint8, mstype.uint16, mstype.uint32, mstype.uint64,
+            mstype.float16, mstype.float64
+        ]
+        if dout_type in type_list:
+            dout = cast_op(dout, mstype.float32)
+        dout_shape = shape_op(dout)
+        axis = tuple([i for i in range(len(dout_shape))])
+        dvalue = sum_op(dout, axis)
+        return zeros_like(shape), cast_op(dvalue, dout_type)
+
+    return bprop
 
 
 @bprop_getters.register(StridedSliceV2)
@@ -69,7 +105,7 @@ def get_bprop_strided_slice_v2(self):
 
     def bprop(x, begin, end, strides, out, dout):
         x_shape = shape_op(x)
-        if -1 in x_shape:
+        if F.is_sequence_value_unknown(x_shape):
             x_shape = dyn_shape_op(x)
         dx = input_grad(x_shape, begin, end, strides, dout)
         dx_all = (dx, zeros_like(begin), zeros_like(end), zeros_like(strides))
@@ -114,7 +150,12 @@ def get_bprop_masked_select(self):
         dinput = mul_op(dout, (1 - mask))
         dvalue = mul_op(dout, mask)
         dinput, dvalue = binop_grad_common(input_data, mask, dinput, dvalue)
-        dvalue = sum_op(dvalue)
+        # for dynamic rank, reduce axis should be calc
+        if F.is_sequence_shape_unknown(P.Shape()(dvalue)):
+            axis = P.Range()(Tensor(0), dyn_rank(dvalue), Tensor(1))
+            dvalue = sum_op(dvalue, axis)
+        else:
+            dvalue = sum_op(dvalue)
         dinput = F.cast(dinput, F.dtype(input_data))
         if is_instance_op(value, mstype.number):
             dvalue = 0
@@ -163,6 +204,16 @@ def get_bprop_masked_scatter(self):
     return bprop
 
 
+@bprop_getters.register(CountNonZero)
+def get_bprop_countnonzero(self):
+    """Grad definition for CountNonZero"""
+
+    def bprop(x, out, dout):
+        return (zeros_like(x),)
+
+    return bprop
+
+
 @bprop_getters.register(Mvlgamma)
 def get_bprop_mvlgamma(self):
     """Grad definition for Mvlgamma"""
@@ -210,7 +261,7 @@ def get_bprop_index_fill(self):
     def bprop(x, dim, indices, value, out, dout):
         zero_value = zeros_like(value)
         x_grad = index_fill(dout, dim, indices, zero_value)
-        if is_shape_unknown(shape(x)):
+        if F.is_sequence_value_unknown(shape(x)):
             if dyn_rank(x) == 0:
                 value_grad = dout
             else:
@@ -226,6 +277,43 @@ def get_bprop_index_fill(self):
     return bprop
 
 
+@bprop_getters.register(IndexPut)
+def get_bprop_index_put(self):
+    """Generate bprop for IndexPut"""
+    gather_nd = P.GatherNd()
+    stack = P.Stack()
+    tile = P.Tile()
+    masked_select = MaskedSelect()
+    masked_scatter = MaskedScatter()
+    accumulate_grad = self.accumulate
+    index_put = IndexPut(accumulate=accumulate_grad)
+    is_ascend = context.get_context("device_target") == 'Ascend'
+
+    # Negative value are not supported for GatherNd indices when Ascend, so convert it to positive value.
+    def convert_idx_positive(indices_i, x_shape_i):
+        mask = indices_i < 0
+        idx_pos = masked_select(indices_i + x_shape_i, mask)
+        idx = masked_scatter(indices_i, mask, idx_pos)
+        return idx
+
+    def bprop(x1, x2, indices, out, dout):
+        maxsize = max(x.shape[0] for x in indices)
+        indices_ms = [tile(x, (maxsize,)) if x.shape[0] == 1 else x for x in indices]
+        if is_ascend:
+            indices_ms = [convert_idx_positive(indices_ms[i], x1.shape[i]) for i in range(len(indices_ms))]
+        indices_grad = stack(indices_ms).T
+        values_grad = gather_nd(dout, indices_grad)
+        if x2.shape[0] == 1:
+            values_grad = values_grad.sum().reshape(1)
+        if values_grad.shape != x2.shape and len(indices) < len(x1.shape):
+            _, values_grad = binop_grad_common(x1, x2, dout, values_grad)
+        if accumulate_grad == 0:
+            dout = index_put(dout, zeros_like(x2), indices)
+        return dout, values_grad, [zeros_like(item) for item in indices]
+
+    return bprop
+
+
 @bprop_getters.register(P.TensorScatterSub)
 def get_bprop_tensor_scatter_sub(self):
     """Generate bprop for TensorScatterSub"""
@@ -280,7 +368,7 @@ def get_bprop_matrix_diag_part_v3(self):
 
     def bprop(x, k, padding_value, out, dout):
         shape_this = P.Shape()(x)[-2:]
-        if not is_shape_unknown(shape_this):
+        if not F.is_sequence_value_unknown(shape_this):
             row = shape_this[0]
             col = shape_this[1]
             result = (matrix_diag_v3(dout, k, Tensor(row, dtype=mstype.int32), Tensor(col, dtype=mstype.int32),
@@ -304,7 +392,7 @@ def get_bprop_matrix_set_diag_v3(self):
         diagonal_cal = matrix_diag_part_v3(dout, k, zeros((), dout.dtype))
 
         diagonal_shape = P.Shape()(diagonal)
-        if is_shape_unknown(diagonal_shape):
+        if F.is_sequence_value_unknown(diagonal_shape):
             diagonal = F.cast(diagonal, dout.dtype)
             x_cal = matrix_set_diag_v3(dout, zeros_like(diagonal), k)
         else:
@@ -327,7 +415,7 @@ def tensor_scatter_possible_replacement(x, indices, updates, out, dout):
     possibly_updated = gather_nd(out, indices)
     out_indicators = F.cast(equal(updates, possibly_updated), mstype.int32)
     input_shape = shape(x)
-    if is_shape_unknown(input_shape):
+    if F.is_sequence_value_unknown(input_shape):
         input_shape = dyn_shape_op(x)
 
     scattered_out_indicators = scatter_nd(indices, out_indicators, input_shape)
@@ -347,6 +435,15 @@ def get_bprop_log_normal_reverse(self):
     return bprop
 
 
+@bprop_getters.register(ParameterizedTruncatedNormal)
+def get_bprop_parameterized_truncated_normal(self):
+    """Grad definition for `ParameterizedTruncatedNormal` operation."""
+    def bprop(shape, mean, stdevs, min_val, max_val, out, dout):
+        return (zeros_like(shape), zeros_like(mean), zeros_like(stdevs), zeros_like(min_val), zeros_like(max_val))
+
+    return bprop
+
+
 @bprop_getters.register(P.TensorScatterMax)
 def get_bprop_tensor_scatter_max(self):
     """Generate bprop for TensorScatterMax"""
@@ -446,13 +543,13 @@ def get_bprop_resize_nearest_neighbor_v2(self):
 
     def bprop(x, size, output, dout):
         x_shape = P.Shape()(x)
-        if is_shape_unknown(x_shape):
+        if F.is_sequence_value_unknown(x_shape):
             x_shape = P.TensorShape()(x)
         grad_in_size = x_shape[1:3]
         if data_format == 'NCHW':
             grad_in_size = x_shape[2:4]
 
-        if is_shape_unknown(P.Shape()(x)):
+        if F.is_sequence_value_unknown(P.Shape()(x)):
             dx = grad_op(dout, grad_in_size)
             return dx, zeros_like(grad_in_size)
 
@@ -469,7 +566,7 @@ def get_bprop_col2im(self):
     dilations = self.dilation
     strides = self.stride
     pads = self.padding
-    im2col = Im2Col(ksizes=ksizes, dilations=dilations, strides=strides, padding_mode="CALCULATED", pads=pads)
+    im2col = Im2Col(ksizes=ksizes, dilations=dilations, strides=strides, pads=pads)
 
     def bprop(x, output_size, out, dout):
         dx = im2col(dout)
@@ -478,6 +575,36 @@ def get_bprop_col2im(self):
     return bprop
 
 
+@bprop_getters.register(Im2Col)
+def get_bprop_im2col(self):
+    """
+    Generate bprop for Im2Col
+
+    Im2Col, corresponding to torch's UnFold operator.
+    The Unfold operator has no `padding_mode` attribute,
+    and it's implementation corresponds to the mindspore
+    implementation with `padding_mode=CALCULATED` .
+    So, currently the bprop function of Im2Col only supports
+    the CALCULATED mode.
+    """
+    kernel_size = self.ksizes
+    dilation = self.dilations
+    stride = self.strides
+    padding = (self.pads[0], self.pads[-1])
+    shape_op = P.TensorShape()
+    col2im = Col2Im(kernel_size=kernel_size,
+                    dilation=dilation,
+                    stride=stride,
+                    padding=padding)
+
+    def bprop(x, out, dout):
+        x_shape = shape_op(x)[2:]
+        dx = col2im(dout, x_shape)
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(P.ExtractVolumePatches)
 def get_bprop_extract_volume_patches(self):
     """Generate bprop for ExtractVolumePatches"""
@@ -538,7 +665,7 @@ def get_bprop_extract_volume_patches(self):
     def bprop(x, out, dout):
         x_shape = P.Shape()(x)
         out_shape = P.Shape()(out)
-        if is_shape_unknown(x_shape) or is_shape_unknown(out_shape):
+        if F.is_sequence_value_unknown(x_shape) or F.is_sequence_value_unknown(out_shape):
             return _dyn_extract_volume_patches(x, out, dout)
         x_n, x_c, x_d, x_h, x_w = x_shape
         x_indices_num = 1 + x_d * x_h * x_w
@@ -609,6 +736,7 @@ def get_bprop_affinegrid(self):
     """Generate bprop for AffineGrid"""
 
     align_corners = self.align_corners
+    input_grad = G.AffineGridGrad(align_corners)
     ones = P.Ones()
     transpose = P.Transpose()
     concat = P.Concat(1)
@@ -824,12 +952,19 @@ def get_bprop_affinegrid(self):
             dtheta = transpose(dtheta, perm2)
         return dtheta, tre
 
-    def bprop(theta, output_size, out, dout):
+    def bprop_gpu(theta, output_size, out, dout):
         is_tensor, _ = convert_to_tensor(output_size)
         if is_tensor:
             return dyn_bprop(theta, output_size, out, dout)
         return static_bprop(theta, output_size, out, dout)
 
+    def bprop(theta, output_size, out, dout):
+        dx = input_grad(dout, output_size)
+        return dx, zeros_like(output_size)
+
+    if context.get_context('device_target') == "GPU":
+        return bprop_gpu
+
     return bprop
 
 
@@ -907,7 +1042,7 @@ def get_bprop_expand(self):
 
     def bprop(x, shape, out, dout):
         reduce_dims = []
-        dshape = zeroslike(dout)
+        dshape = zeroslike(shape)
         dx_shape = dout.shape
         if dx_shape is None:
             return dout.sum(), dshape
@@ -947,12 +1082,12 @@ def get_bprop_segment_mean(self):
         dout_type = F.dtype(dout)
 
         ones_shape = shape(segment_ids)
-        if is_shape_unknown(ones_shape):
+        if F.is_sequence_value_unknown(ones_shape):
             ones_shape = dyn_shape(segment_ids)
 
         ones = ()
         inputx_shape = shape(input_x)
-        if is_shape_unknown(inputx_shape):
+        if F.is_sequence_value_unknown(inputx_shape):
             input_rank = dyn_rank(input_x)
             if input_rank > cast(1, mstype.float32):
                 ones_shape = concat([ones_shape, dyn_ones(expand_dims(input_rank - 1, 0), mstype.int64)])