mindspore 2.0.0a0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/ops/_grad/grad_math_ops.py

@@ -15,7 +15,6 @@
 
 """Define the grad rules of math related operations."""
 
-from functools import reduce
 import numpy as np
 import mindspore as ms
 from mindspore import nn
@@ -30,10 +29,9 @@ from mindspore.ops._grad.grad_base import bprop_getters, create_tensor_by_elemen
 from mindspore.ops._grad.grad_base import convert_to_tensor
 from mindspore.ops._grad.grad_base import sum_grad_reduce_axis, dyn_fill, dyn_rank
 from mindspore.ops._grad.grad_base import dyn_ones, dyn_rank_1d
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops.composite.multitype_ops import _constexpr_utils as const_utils
-from mindspore.ops.operations._inner_ops import DynamicBroadcastGradientArgs, DynamicBroadcastTo, IsSubClass
-from mindspore.ops._utils.utils import is_shape_unknown, is_dim_unknown
+from mindspore.ops.operations._inner_ops import DynamicBroadcastGradientArgs, IsSubClass, DynamicBroadcastTo
 from mindspore.ops.operations import array_ops as A
 
 shape_op = P.Shape()
@@ -114,7 +112,7 @@ def binop_grad_common(x, y, dx, dy):
     # if input shape is the same as dout shape, do not need to reduce
     reduce_dx = dx
     reduce_dy = dy
-    if not (is_shape_unknown(shape_of_x) or is_shape_unknown(shape_of_y)):
+    if not (F.is_sequence_value_unknown(shape_of_x) or F.is_sequence_value_unknown(shape_of_y)):
         rx = broadcast_gradient_args(shape_of_x, shape_of_y)
         if rx[0]:
             # if dx is scalar whose shape is (), do not need reduce
@@ -127,11 +125,12 @@ def binop_grad_common(x, y, dx, dy):
                 dy = _reduce_sum_with_cast(dy, rx[1])
             reduce_dy = reshape(dy, shape_of_y)
         return reduce_dx, reduce_dy
-    if not shape_of_x or not shape_of_y:
+
+    if not isinstance(shape_of_x, tuple) or not isinstance(shape_of_y, tuple):
         # x or y is scalar
-        if not shape_of_x:
+        if not isinstance(shape_of_x, tuple):
             reduce_dx = _reduce_sum_with_cast(dx, ())
-        if not shape_of_y:
+        if not isinstance(shape_of_y, tuple):
             reduce_dy = _reduce_sum_with_cast(dy, ())
         return reduce_dx, reduce_dy
 
@@ -151,7 +150,7 @@ def binop_grad_common_with_shift(x, y, dx, dy, shift):
     # if input shape is the same as dout shape, do not need to reduce
     reduce_dx = dx
     reduce_dy = dy
-    if not (is_shape_unknown(broadcast_shape_of_x) or is_shape_unknown(broadcast_shape_of_y)):
+    if not (F.is_sequence_value_unknown(broadcast_shape_of_x) or F.is_sequence_value_unknown(broadcast_shape_of_y)):
         rx = broadcast_gradient_args(broadcast_shape_of_x, broadcast_shape_of_y)
         if rx[0]:
             # if dx is scalar whose shape is (), do not need reduce
@@ -164,11 +163,12 @@ def binop_grad_common_with_shift(x, y, dx, dy, shift):
                 dy = _reduce_sum_with_cast(dy, rx[1])
             reduce_dy = reshape(dy, shape_of_y)
         return reduce_dx, reduce_dy
-    if not shape_of_x or not shape_of_y:
+
+    if not isinstance(shape_of_x, tuple) or not isinstance(shape_of_y, tuple):
         # x or y is scalar
-        if not shape_of_x:
+        if not isinstance(shape_of_x, tuple):
             reduce_dx = _reduce_sum_with_cast(dx, ())
-        if not shape_of_y:
+        if not isinstance(shape_of_y, tuple):
             reduce_dy = _reduce_sum_with_cast(dy, ())
         return reduce_dx, reduce_dy
 
@@ -178,7 +178,7 @@ def binop_grad_common_with_shift(x, y, dx, dy, shift):
 def _dyn_reduced_shape(input_shape, axis, x):
     """Dynamic reduce shape"""
     input_shape = P.Cast()(input_shape, ms.int32)
-    if x is not None and not is_dim_unknown(shape_op(x)):
+    if x is not None and not F.is_sequence_shape_unknown(shape_op(x)):
         input_rank = len(shape_op(x))
     else:
         input_rank = dyn_rank(x)
@@ -209,7 +209,7 @@ def _sum_grad(x, axis, dout):
     """Grad definition for `Sum` operation."""
     input_shape = shape_op(x)
     is_mutable, axis = convert_to_tensor(axis)
-    if is_shape_unknown(input_shape) or is_mutable:
+    if F.is_sequence_value_unknown(input_shape) or is_mutable:
         input_shape = dyn_shape_op(x)
         output_shape_kept_dims = _dyn_reduced_shape(input_shape, axis, x)
         output_shape_kept_dims = P.Cast()(output_shape_kept_dims, ms.int32)
@@ -226,7 +226,7 @@ def _min_or_max_grad(x, axis, out, dout):
     """Grad definition for `Min` and `Max` operations."""
     input_shape = shape_op(x)
     output_shape_kept_dims = ()
-    if is_shape_unknown(input_shape):
+    if F.is_sequence_value_unknown(input_shape):
         input_shape = dyn_shape_op(x)
         output_shape_kept_dims = _dyn_reduced_shape(input_shape, axis, x)
         output_shape_kept_dims = P.Cast()(output_shape_kept_dims, ms.int32)
@@ -268,7 +268,7 @@ def _argmin_or_argmax_grad(x, axis, keep_dims, op, out, dout):
     x_axis = axis
     onehot_axis_is_neg = False
     if x_axis < 0:
-        if not is_dim_unknown(x_shape):
+        if not F.is_sequence_shape_unknown(x_shape):
             x_axis = axis + x_dim
         else:
             onehot_axis_is_neg = True
@@ -279,13 +279,13 @@ def _argmin_or_argmax_grad(x, axis, keep_dims, op, out, dout):
     else:
         dout_expand = expand(dout[1], onehot_axis)
     out_shape = shape_op(out[0])
-    if not is_dim_unknown(out_shape):
+    if not F.is_sequence_shape_unknown(out_shape):
         if onehot_axis >= len(out_shape):
             onehot_axis = -1
     type_x = F.dtype(x)
     on_value = F.cast(F.scalar_to_tensor(1.0), type_x)
     off_value = F.cast(F.scalar_to_tensor(0.0), type_x)
-    if not is_shape_unknown(x_shape):
+    if not F.is_sequence_value_unknown(x_shape):
         depth = 1
         if x_shape:
             depth = x_shape[axis]
@@ -308,35 +308,6 @@ def _argmin_or_argmax_grad(x, axis, keep_dims, op, out, dout):
     return dx
 
 
-@bprop_getters.register(P.MatMul)
-def bprop_matmul(self):
-    """Grad definition for `MatMul` operation."""
-    ta = self.transpose_a
-    tb = self.transpose_b
-    mul1 = P.MatMul(transpose_a=(ta and tb),
-                    transpose_b=(ta or (not tb)))
-    mul2 = P.MatMul(transpose_a=((not ta) or tb),
-                    transpose_b=(ta and tb))
-
-    def bprop(x, w, out, dout):
-        conj = P.Conj()
-        origin_dtype = x.dtype
-        if origin_dtype in (mstype.complex64, mstype.complex128):
-            x = conj(x)
-            w = conj(w)
-        if ta:
-            dx = mul1(w, dout)
-        else:
-            dx = mul1(dout, w)
-        if tb:
-            dw = mul2(dout, x)
-        else:
-            dw = mul2(x, dout)
-        return dx, dw
-
-    return bprop
-
-
 @bprop_getters.register(P.BatchMatMul)
 def bprop_batchmatmul(self):
     """Grad definition for `BatchMatMul` operation."""
@@ -361,16 +332,6 @@ def bprop_batchmatmul(self):
     return bprop
 
 
-@bprop_getters.register(P.Add)
-def get_bprop_add(self):
-    """Grad definition for `Add` operation."""
-
-    def bprop(x, y, out, dout):
-        return binop_grad_common(x, y, dout, dout)
-
-    return bprop
-
-
 @bprop_getters.register(P.TensorAdd)
 def get_bprop_tensor_add(self):
     """Grad definition for `Add` operation."""
@@ -397,35 +358,14 @@ def get_bprop_matrix_inverse(self):
     return bprop
 
 
-@bprop_getters.register(P.Neg)
-def get_bprop_neg(self):
-    """Grad definition for `Neg` operation."""
-    neg_grad = P.Neg()
-
-    def bprop(x, out, dout):
-        dx = neg_grad(dout)
-        return (dx,)
-
-    return bprop
-
-
-@bprop_getters.register(P.Sub)
-def get_bprop_sub(self):
-    """Grad definition for `Sub` operation."""
-    neg_func = P.Neg()
-
-    def bprop(x, y, out, dout):
-        return binop_grad_common(x, y, dout, neg_func(dout))
-
-    return bprop
-
-
 @bprop_getters.register(P.Mul)
 def get_bprop_mul(self):
     """Grad definition for `Mul` operation."""
     mul_func = P.Mul()
 
     def bprop(x, y, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("For 'Mul', gradient not support for complex type currently.")
         bc_dx = mul_func(y, dout)
         bc_dy = mul_func(x, dout)
         return binop_grad_common(x, y, bc_dx, bc_dy)
@@ -441,6 +381,8 @@ def get_bprop_real_div(self):
     mul_op = P.Mul()
 
     def bprop(x, y, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("For 'RealDiv', gradient not support for complex type currently.")
         bc_x = div_op(dout, y)
         bc_y = neg(mul_op(bc_x, out))
         return binop_grad_common(x, y, bc_x, bc_y)
@@ -501,7 +443,7 @@ def get_bprop_floor(self):
     dtype_ = P.DType()
 
     def bprop(x, out, dout):
-        if is_shape_unknown(shape_(x)):
+        if F.is_sequence_value_unknown(shape_(x)):
             bc_x = zeros_like(x)
         else:
             bc_x = fill_(dtype_(x), shape_(x), 0.)
@@ -518,7 +460,7 @@ def get_bprop_ceil(self):
     dtype_ = P.DType()
 
     def bprop(x, out, dout):
-        if is_shape_unknown(shape_(x)):
+        if F.is_sequence_value_unknown(shape_(x)):
             bc_x = zeros_like(x)
         else:
             bc_x = fill_(dtype_(x), shape_(x), 0.)
@@ -537,6 +479,36 @@ def get_bprop_floordiv(self):
     return bprop
 
 
+@bprop_getters.register(P.BitwiseAnd)
+def get_bprop_bitwiseand(self):
+    """Grad definition for `BitwiseAnd` operation."""
+
+    def bprop(x, y, out, dout):
+        return zeros_like(x), zeros_like(y)
+
+    return bprop
+
+
+@bprop_getters.register(P.BitwiseOr)
+def get_bprop_bitwiseor(self):
+    """Grad definition for `BitwiseOr` operation."""
+
+    def bprop(x, y, out, dout):
+        return zeros_like(x), zeros_like(y)
+
+    return bprop
+
+
+@bprop_getters.register(P.BitwiseXor)
+def get_bprop_bitwisexor(self):
+    """Grad definition for `BitwiseXor` operation."""
+
+    def bprop(x, y, out, dout):
+        return zeros_like(x), zeros_like(y)
+
+    return bprop
+
+
 @bprop_getters.register(P.FloorMod)
 def get_bprop_floormod(self):
     """Grad definition for `FloorMod` operation."""
@@ -594,7 +566,7 @@ def get_bprop_square(self):
     def bprop(x, out, dout):
         temp = mul_func(dout, x)
         shape_x = shape_op(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             fill_value = dyn_fill(dtype(temp), dyn_shape_op(x), 2.0)
         else:
             fill_value = fill_func(dtype(temp), shape_x, 2.0)
@@ -644,12 +616,12 @@ def get_bprop_square_sum_all(self):
     def bprop(x, y, out, dout):
         temp_x = mul_func(dout[0], x)
         temp_y = mul_func(dout[1], y)
-        if is_shape_unknown(shape_op(x)):
+        if F.is_sequence_value_unknown(shape_op(x)):
             dx = mul_func(dyn_fill(dtype(temp_x), dyn_shape_op(x), 2.0), temp_x)
         else:
             dx = mul_func(fill_func(dtype(temp_x), shape_op(x), 2.0), temp_x)
 
-        if is_shape_unknown(shape_op(y)):
+        if F.is_sequence_value_unknown(shape_op(y)):
             dy = mul_func(dyn_fill(dtype(temp_y), dyn_shape_op(y), 2.0), temp_y)
         else:
             dy = mul_func(fill_func(dtype(temp_y), shape_op(y), 2.0), temp_y)
@@ -792,9 +764,11 @@ def get_bprop_pow(self):
     ln = P.Log()
 
     def bprop(x, power, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("For 'Pow', gradient not support for complex type currently.")
         bc_dx = power * pow_op(x, power - 1.0) * dout
         shape_x = shape_op(x)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x):
             x = F.select(x < 0, dyn_fill(F.dtype(x), dyn_shape_op(x), 1), x)
         else:
             x = F.select(x < 0, F.fill(F.dtype(x), F.shape(x), 1), x)
@@ -888,21 +862,31 @@ def get_bprop_cumsum(self):
     return bprop
 
 
-@constexpr
+@_primexpr
 def _split_shape_index(input_shape, axis):
     """Calculate reduce_prod grad transpose indices and perm shape."""
     rank = len(input_shape)
     if isinstance(axis, int):
         axis = tuple([axis])
     reduction_indices = tuple([(i + rank) % rank for i in axis])
-    other_indices = tuple(set(range(rank)) - set(reduction_indices))
-    reduced_num = reduce(lambda x, y: x * y, [1] + [input_shape[i] for i in reduction_indices])
-    other_num = reduce(lambda x, y: x * y, [1] + [input_shape[i] for i in other_indices])
+    other_indices_list = []
+    for i in range(rank):
+        if i not in reduction_indices and i not in other_indices_list:
+            other_indices_list.append(i)
+    other_indices = tuple(other_indices_list)
+    reduced_list = [1] + [input_shape[i] for i in reduction_indices]
+    other_list = [1] + [input_shape[i] for i in other_indices]
+    reduced_num = 1
+    for i in reduced_list:
+        reduced_num = reduced_num * i
+    other_num = 1
+    for i in other_list:
+        other_num = other_num * i
     perm = reduction_indices + other_indices
     return tuple([reduced_num, other_num]), perm
 
 
-@constexpr
+@_primexpr
 def _invert_permutation(perm):
     """Calculate invert permutation."""
     out = [0] * len(perm)
@@ -940,9 +924,15 @@ def get_bprop_reduceprod(self):
 
     def bprop(x, axis, out, dout):
         """Grad definition for `Product` operation."""
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("The 'ReduceProd', gradient not support for complex type currently.")
         # Expand dout to full input shape
         input_shape = shape_op(x)
-        if is_shape_unknown(input_shape):
+        if input_shape == ():
+            dx = _sum_grad(x, axis, dout)
+            return dx, zeros_like(axis)
+
+        if F.is_sequence_value_unknown(input_shape):
             input_shape = dyn_shape_op(x)
             input_shape = P.Cast()(input_shape, ms.int64)
             output_shape_kept_dims = _dyn_reduced_shape(input_shape, axis, x)
@@ -953,14 +943,14 @@ def get_bprop_reduceprod(self):
         dout = reshape(dout, output_shape_kept_dims)
 
         # Pack all reduced dimensions into a single one, so we can perform the cumprod ops.
-        if is_shape_unknown(shape_op(x)):
+        if F.is_sequence_value_unknown(shape_op(x)):
             pack_shape, perm = _split_dyn_shape_index(x, axis)
         else:
             pack_shape, perm = _split_shape_index(shape_op(x), axis)
 
         permuted = transpose(x, perm)
         permuted_shape = shape_op(permuted)
-        if is_shape_unknown(permuted_shape):
+        if F.is_sequence_value_unknown(permuted_shape):
             permuted_shape = dyn_shape_op(permuted)
             pack_shape = create_tensor_by_element(pack_shape)
         reshaped = reshape(permuted, pack_shape)
@@ -972,7 +962,7 @@ def get_bprop_reduceprod(self):
 
         # Invert the transpose and reshape operations.
         # Make sure to set the statically known shape information through a reshape.
-        if is_shape_unknown(shape_op(permuted)):
+        if F.is_sequence_value_unknown(shape_op(permuted)):
             dout = DynamicBroadcastTo()(dout, input_shape)
             out = transpose(y, dyn_invert_permutation(perm)) * dout
         else:
@@ -1027,6 +1017,8 @@ def get_bprop_reducemax(self):
     """Grad definition for `Max` operation."""
 
     def bprop(x, axis, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("The 'ReduceMax', gradient not support for complex type currently.")
         dx = _min_or_max_grad(x, axis, out, dout)
         return (dx, zeros_like(axis))
 
@@ -1052,6 +1044,8 @@ def get_bprop_reducemin(self):
     """Grad definition for `ReduceMin` operation."""
 
     def bprop(x, axis, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("The 'ReduceMin', gradient not support for complex type currently.")
         dx = _min_or_max_grad(x, axis, out, dout)
         return (dx, zeros_like(axis))
 
@@ -1080,10 +1074,12 @@ def get_bprop_reduce_mean(self):
     dtype = P.DType()
 
     def bprop(x, axis, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("The 'ReduceMean', gradient not support for complex type currently.")
         grad = _sum_grad(x, axis, dout)
         shape_x = shape_op(x)
         shape_out = shape_op(out)
-        if is_shape_unknown(shape_x):
+        if F.is_sequence_value_unknown(shape_x) or F.is_sequence_value_unknown(shape_out):
             shape_x = dyn_shape_op(x)
             shape_out = dyn_shape_op(out)
             div_shape = reduce_prod(cast(shape_x, mstype.float32), ()) /\
@@ -1091,7 +1087,7 @@ def get_bprop_reduce_mean(self):
             dx = div_op(grad, cast(div_shape, dtype(grad)))
         else:
             div_shape = F.shape_mul(shape_x) / F.shape_mul(shape_out)
-            dx = div_op(grad, cast(F.scalar_to_tensor(div_shape, mstype.int32), dtype(grad)))
+            dx = div_op(grad, cast(F.scalar_to_tensor(div_shape), dtype(grad)))
         return dx, zeros_like(axis)
 
     return bprop
@@ -1217,16 +1213,6 @@ def get_bprop_logical_and(self):
     return bprop
 
 
-@bprop_getters.register(P.LogicalOr)
-def get_bprop_logical_or(self):
-    """Grad definition for `LogicalOr` operation."""
-
-    def bprop(x, y, out, dout):
-        return zeros_like(x), zeros_like(y)
-
-    return bprop
-
-
 @bprop_getters.register(P.NPUAllocFloatStatus)
 def get_bprop_npu_alloc_float_status(self):
     """Grad definition for `NPUAllocFloatStatus` operation."""
@@ -1424,6 +1410,9 @@ def get_bprop_cosh(self):
     sinh = P.Sinh()
 
     def bprop(x, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("The 'Cosh', gradient not support for complex type currently.")
+
         dx = sinh(x) * dout
         return (dx,)
 
@@ -1454,16 +1443,6 @@ def get_bprop_conj(self):
     return bprop
 
 
-@bprop_getters.register(P.ScalarCast)
-def get_bprop_scalar_cast(self):
-    """Generate bprop for ScalarCast"""
-
-    def bprop(x, t, out, dout):
-        return F.scalar_cast(dout, F.typeof(x)), zeros_like(t)
-
-    return bprop
-
-
 @bprop_getters.register(P.AccumulateNV2)
 def get_bprop_scalar_accumulatenv2(self):
     """Generate bprop for AccumulateNV2"""
@@ -1577,6 +1556,9 @@ def get_bprop_tan(self):
     cos = P.Cos()
 
     def bprop(x, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("For 'Tan', gradient not support for complex type currently.")
+
         cosx = cos(x)
         secx2 = square(reciprocal(cosx))
         dx = secx2 * dout
@@ -1618,6 +1600,9 @@ def get_bprop_atanh(self):
     div = P.Div()
 
     def bprop(x, out, dout):
+        if x.dtype in (mstype.complex64, mstype.complex128):
+            raise TypeError("For 'Atanh', gradient not support for complex type currently.")
+
         tmp = 1 - power(x, 2)
         dx = div(1, tmp) * dout
         return (dx,)
@@ -1657,3 +1642,43 @@ def get_bprop_index_add(self):
         return dout, zeros_like(indices), gather(dout, indices, _axis)
 
     return bprop
+
+
+@bprop_getters.register(P.InplaceUpdate)
+def get_bprop_inplace_update(self):
+    """Grad definition for `InplaceUpdate` operation."""
+
+    def bprop(x, v, out, dout):
+        return zeros_like(x), zeros_like(v)
+
+    return bprop
+
+
+@bprop_getters.register(P.InplaceUpdateV2)
+def get_bprop_inplace_update_v2(self):
+    """Grad definition for `InplaceUpdateV2` operation."""
+
+    def bprop(x, indices, v, out, dout):
+        return zeros_like(x), zeros_like(indices), zeros_like(v)
+
+    return bprop
+
+
+@bprop_getters.register(P.InplaceSub)
+def get_bprop_inplace_sub(self):
+    """Grad definition for `InplaceSub` operation."""
+
+    def bprop(x, input_v, out, dout):
+        return zeros_like(x), zeros_like(input_v)
+
+    return bprop
+
+
+@bprop_getters.register(P.InplaceAdd)
+def get_bprop_inplace_add(self):
+    """Grad definition for `InplaceAdd` operation."""
+
+    def bprop(x, input_v, out, dout):
+        return zeros_like(x), zeros_like(input_v)
+
+    return bprop