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

mindspore/ops/_grad/grad_nn_ops.py

@@ -14,11 +14,10 @@
 # ============================================================================
 
 """Define the grad rules of neural network related operations."""
-import numpy as np
 from mindspore import context
 from mindspore.common import dtype as mstype
 from mindspore.common.tensor import Tensor
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops.operations import nn_ops as nps
 from mindspore.ops._grad.grad_base import bprop_getters, dyn_size, create_tensor_by_element, dyn_rank
 from mindspore.ops import functional as F
@@ -27,38 +26,29 @@ from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.ops.operations import _rl_inner_ops as rl_ops
-from mindspore.ops._utils.utils import range_op, get_1d_shape, is_shape_unknown, is_dim_unknown
+from mindspore.ops._utils.utils import range_op, get_1d_shape
 
 
-@bprop_getters.register(P.BiasAdd)
-def get_bprop_bias_add(self):
-    """Grad definition for `BiasAdd` operation."""
-    bias_grad = G.BiasAddGrad(self.data_format)
-
-    def bprop(x, w, out, dout):
-        return dout, bias_grad(dout)
+@_primexpr
+def bias_add_gradgrad_helper(shape, bias_shape, data_format):
+    """Helper function of BiasGradGrad to calculate expanded shape."""
+    new_shape = list(shape)
+    new_bias_shape = list(bias_shape)
 
-    return bprop
+    ones_1 = []
+    ones_2 = []
+    for _ in new_shape[2:]:
+        ones_1.append(1)
 
+    for _ in new_shape[:-1]:
+        ones_2.append(1)
 
-@constexpr
-def bias_add_gradgrad_helper(shape, bias_shape, data_format):
-    """Helper function of BiasGradGrad to calculate expanded shape."""
-    shape = np.array(shape).astype(np.int)
-    bias_shape = np.array(bias_shape).astype(np.int)
     if data_format == "NCHW":
-        expanded_shape = np.concatenate([
-            np.ones_like(shape[:1]),
-            bias_shape,
-            np.ones_like(shape[2:])
-        ], axis=0)
-        tile_mults = np.concatenate([shape[:1], [1], shape[2:]], axis=0)
+        expanded_shape = [1] + new_bias_shape + ones_1
+        tile_mults = [new_shape[0]] + [1] + new_shape[2:]
     else:
-        expanded_shape = np.concatenate([
-            np.ones_like(shape[:-1]),
-            bias_shape
-        ], axis=0)
-        tile_mults = np.concatenate([shape[:-1], [1]], axis=0)
+        expanded_shape = ones_2 + new_bias_shape
+        tile_mults = new_shape[:-1] + [1]
     return tuple(expanded_shape), tuple(tile_mults)
 
 
@@ -85,7 +75,7 @@ def get_bprop_bias_add_grad(self):
         dyn_shape = P.TensorShape()
         dy_shape = dy.shape
         dout_shape = dout.shape
-        if is_shape_unknown(dy_shape) or is_shape_unknown(dout_shape):
+        if F.is_sequence_value_unknown(dy_shape) or F.is_sequence_value_unknown(dout_shape):
             dy_shape = dyn_shape(dy)
             dout_shape = dyn_shape(dout)
             expanded_shape, tile_mults = bias_add_gradgrad_helper_dynamic(dy_shape, dout_shape, data_format)
@@ -100,38 +90,6 @@ def get_bprop_bias_add_grad(self):
     return bprop
 
 
-@bprop_getters.register(P.Conv2D)
-def get_bprop_conv2d(self):
-    """Grad definition for `Conv2D` operation."""
-    self.out_channel = self.get_attr_dict()["out_channel"]
-    self.pad_list = self.get_attr_dict()["pad_list"]
-    input_grad = P.Conv2DBackpropInput(
-        self.out_channel, self.kernel_size, self.pad_mode, self.pad, self.pad_list, mode=self.mode,
-        dilation=self.dilation, stride=self.stride, group=self.group, data_format=self.format
-    )
-    filter_grad = G.Conv2DBackpropFilter(
-        self.out_channel, self.kernel_size, self.pad_mode, self.pad, self.pad_list, mode=self.mode,
-        dilation=self.dilation, stride=self.stride, group=self.group, data_format=self.format
-    )
-    get_shape = P.Shape()
-    get_dyn_shape = P.TensorShape()
-    cast_type = P.Cast()
-
-    def bprop(x, w, out, dout):
-        x_shape = get_shape(x)
-        w_shape = get_shape(w)
-        if is_shape_unknown(x_shape):
-            x_shape = get_dyn_shape(x)
-        if is_shape_unknown(w_shape):
-            w_shape = get_dyn_shape(w)
-            w_shape = cast_type(w_shape, mstype.int32)
-        dx = input_grad(dout, w, x_shape)
-        dw = filter_grad(dout, x, w_shape)
-        return dx, dw
-
-    return bprop
-
-
 @bprop_getters.register(nps.Conv3D)
 def get_bprop_conv3d(self):
     """Grad definition for `Conv3D` operation."""
@@ -145,15 +103,17 @@ def get_bprop_conv3d(self):
     )
     get_shape = P.Shape()
     get_dyn_shape = P.TensorShape()
+    cast = P.Cast()
+    get_dtype = P.DType()
 
     def bprop(x, w, out, dout):
-        if is_shape_unknown(get_shape(x)) or is_shape_unknown(get_shape(w)):
+        if F.is_sequence_value_unknown(get_shape(x)) or F.is_sequence_value_unknown(get_shape(w)):
             dx = input_grad(w, dout, get_dyn_shape(x))
-            dw = filter_grad(x, dout, get_dyn_shape(w))
+            dw = cast(filter_grad(x, dout, get_dyn_shape(w)), get_dtype(x))
             return dx, dw
 
         dx = input_grad(w, dout, get_shape(x))
-        dw = filter_grad(x, dout, get_shape(w))
+        dw = cast(filter_grad(x, dout, get_shape(w)), get_dtype(x))
         return dx, dw
 
     return bprop
@@ -176,7 +136,7 @@ def get_bprop_conv3d_transpose(self):
     get_dyn_shape = P.TensorShape()
 
     def bprop(x, w, out, dout):
-        if is_shape_unknown(F.shape(w)):
+        if F.is_sequence_value_unknown(F.shape(w)):
             dx = input_grad(dout, w)
             dw = filter_grad(dout, x, get_dyn_shape(w))
             return dx, dw
@@ -249,7 +209,7 @@ def get_bprop_extract_image_patches(self):
     def bprop(x, out, dout):
         x_shape = get_shape(x)
         out_shape = get_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_image_patched(x, out, dout)
         x_batch, x_depth, x_row, x_col = x_shape
         x_indices_num = x_row * x_col + 1
@@ -348,6 +308,7 @@ def get_bprop_max_pool_grad_grad(self):
     op_range = P.Range()
     dyn_broadcast_op = inner.DynamicBroadcastTo()
 
+
     def bprop(x1, x2, grad, out, dout):
         dx1 = zeros_like(x1)
         dx2 = zeros_like(x2)
@@ -355,7 +316,7 @@ def get_bprop_max_pool_grad_grad(self):
             dgrad = maxpool_grad_grad(x1, x2, dout)
         else:
             shape_x2 = shape_op(x2)
-            if is_shape_unknown(shape_x2):
+            if F.is_sequence_value_unknown(shape_x2):
                 shape_x2 = dyn_shape_op(x2)
                 b, c, h, w = shape_x2
                 _, ind = maxpool_with_argmax(x1)
@@ -394,22 +355,6 @@ def get_bprop_max_pool_grad_grad_grad(self):
     return bprop
 
 
-@bprop_getters.register(P.MaxPool)
-def get_bprop_max_pool_grad(self):
-    """Grad definition for `MaxPool` operation."""
-    maxpool_grad = G.MaxPoolGrad(
-        kernel_size=self.kernel_size,
-        strides=self.strides,
-        pad_mode=self.pad_mode,
-        data_format=self.format)
-
-    def bprop(x, out, dout):
-        dx = maxpool_grad(x, out, dout)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(P.MaxPool3D)
 def get_bprop_max_pool3d_grad(self):
     """Grad definition for `MaxPool3D` operation."""
@@ -494,9 +439,10 @@ def get_bprop_avg_pool_grad(self):
 def get_bprop_adaptive_avg_pool2d_grad(self):
     """Grad definition for `AdaptiveAvgPool2D` operation."""
     adaptive_avgpool_grad = G.AdaptiveAvgPool2DGrad()
+    shape = P.TensorShape()
 
     def bprop(x, out, dout):
-        dx = adaptive_avgpool_grad(x, dout)
+        dx = adaptive_avgpool_grad(dout, shape(x))
         return (dx,)
 
     return bprop
@@ -518,7 +464,7 @@ def get_bprop_avg_pool_3d_grad(self):
 
     def bprop(x, out, dout):
         x_shape = F.shape(x)
-        if is_shape_unknown(x_shape):
+        if F.is_sequence_value_unknown(x_shape):
             x_shape = P.TensorShape()(x)
         dx = avgpool3d_grad(x_shape, dout)
         return (dx,)
@@ -599,18 +545,6 @@ def get_bprop_mul_no_nan(self):
     return bprop
 
 
-@bprop_getters.register(P.ReLU)
-def get_bprop_relu(self):
-    """Grad definition for `ReLU` operation."""
-    input_grad = G.ReluGrad()
-
-    def bprop(x, out, dout):
-        dx = input_grad(dout, out)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(G.ReluGrad)
 def get_bprop_relu_grad(self):
     """Grad definition for `ReLUGrad` operation."""
@@ -709,7 +643,7 @@ def get_bprop_sigmoid_grad(self):
     return bprop
 
 
-@constexpr
+@_primexpr
 def _get_transpose_axis(x_shp, axis):
     rank = len(x_shp)
     if axis < 0:
@@ -722,7 +656,7 @@ def _get_transpose_axis(x_shp, axis):
 
 def _get_dyn_transpose_axis(x, axis, is_ascend):
     """Get transpose axis"""
-    if is_dim_unknown(P.Shape()(x)):
+    if F.is_sequence_shape_unknown(P.Shape()(x)):
         rank = dyn_rank(x)
         start = Tensor(0, dtype=mstype.int64)
         delta = Tensor(1, dtype=mstype.int64)
@@ -764,13 +698,13 @@ def get_bprop_softmax(self):
     is_ascend = (device_target == "Ascend")
 
     def bprop(x, out, dout):
-        # dx = (dout - sum(dout * out)) * out
+        # dx can be expressed as (dout - sum(dout * out)) * out
         # This formula is correct only when the `axis` is the last dimension.
         # In order to support the scenario where the `axis` is other values,
         # we transpose the data of the `axis` dimension to the last dimension for calculation,
         # and then transpose it back after the calculation.
         shp = get_shape(x)
-        if is_shape_unknown(shp):
+        if F.is_sequence_value_unknown(shp):
             reverse_axis = _get_dyn_transpose_axis(x, axis, is_ascend)
             if is_ascend:
                 reverse_axis = P.Cast()(reverse_axis, mstype.int32)
@@ -855,19 +789,6 @@ def get_bprop_tanh_grad(self):
     return bprop
 
 
-@bprop_getters.register(P.Gelu)
-@bprop_getters.register(P.GeLU)
-def get_bprop_gelu(self):
-    """Grad definition for `GeLU` operation."""
-    input_grad = G.GeLUGrad()
-
-    def bprop(x, out, dout):
-        dx = input_grad(dout, x, out)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(P.FastGeLU)
 def get_bprop_fast_gelu(self):
     """Grad definition for `FastGeLU` operation."""
@@ -909,30 +830,6 @@ def get_bprop_instance_norm(self):
     return bprop
 
 
-@bprop_getters.register(P.BatchNorm)
-def get_bprop_batch_norm(self):
-    """Grad definition for `BatchNorm` operation."""
-    is_training = self.is_training
-    input_grad = G.BatchNormGrad(is_training, self.epsilon, self.data_format)
-
-    def bprop(x, scale, b, mean, variance, out, dout):
-        if is_training:
-            saved_mean = out[3]
-            saved_variance = out[4]
-            reserve = out[2]
-        else:
-            saved_mean = mean
-            saved_variance = variance
-            reserve = out[2]
-        out = input_grad(dout[0], x, scale, saved_mean, saved_variance, reserve)
-        dx = out[0]
-        dscale = out[1]
-        dbias = out[2]
-        return dx, dscale, dbias, zeros_like(mean), zeros_like(variance)
-
-    return bprop
-
-
 @bprop_getters.register(G.BatchNormGrad)
 def get_bprop_batch_norm_grad(self):
     """Grad definition for `BatchNorm` operation."""
@@ -945,19 +842,6 @@ def get_bprop_batch_norm_grad(self):
     return bprop
 
 
-@bprop_getters.register(P.LayerNorm)
-def get_bprop_layer_norm(self):
-    """Grad definition for `LayerNorm` operation."""
-    layer_norm_grad = G.LayerNormGrad(self.begin_norm_axis, self.begin_params_axis)
-
-    def bprop(x, gamma, beta, out, dout):
-        dx, d_gamma, d_beta = layer_norm_grad(
-            x, dout[0], out[2], out[1], gamma)
-        return dx, d_gamma, d_beta
-
-    return bprop
-
-
 @bprop_getters.register(G.LayerNormGrad)
 def get_bprop_layer_norm_grad(self):
     """Grad definition for `LayerNormGrad` operation."""
@@ -1072,7 +956,7 @@ def get_bprop_top_kv2(self):
         ind_2d = reshape_op(indices, (-1, ind_lastdim))
         outerdim = shape_op(ind_2d)[0]
 
-        # [0, outterdim, 2*outerdim, ..., (k-1)*outerdim]
+        # range_flatten_index can be expressed as: [0, outterdim, 2*outerdim, ..., (k-1)*outerdim]
         indices_dtype = dtype(indices)
         range_flatten_index = range_op(0, outerdim * in_lastdim, in_lastdim, indices_dtype)
 
@@ -1099,7 +983,7 @@ def get_bprop_top_kv2(self):
         ind_2d = reshape_op(indices, create_tensor_by_element((-1, ind_lastdim)))
         outerdim = dyn_shape(ind_2d)[0]
 
-        # [0, outterdim, 2*outerdim, ..., (k-1)*outerdim]
+        # range_flatten_index can be expressed as: [0, outterdim, 2*outerdim, ..., (k-1)*outerdim]
         range_flatten_index = P.Range()(cast(0, mstype.int64), outerdim * in_lastdim, in_lastdim)
 
         # expand_dims to (k, 1), then broadcast
@@ -1115,7 +999,7 @@ def get_bprop_top_kv2(self):
         return out_grad, zeros_like(k)
 
     def bprop(input_x, k, out, dout):
-        if is_shape_unknown(shape_op(input_x)):
+        if F.is_sequence_value_unknown(shape_op(input_x)):
             return _bprop_dynshape(input_x, k, out, dout)
         return _bprop_static(input_x, k, out, dout)
 
@@ -1219,6 +1103,27 @@ def get_bprop_lstm(self):
     return bprop
 
 
+@bprop_getters.register(rl_ops.GRUV2)
+def get_bppro_gru_v2(self):
+    """Grad definition for `GRUV2` operation."""
+    gru_grad_v2 = G.GRUV2Grad(
+        self.input_size,
+        self.hidden_size,
+        self.num_layers,
+        self.has_bias,
+        self.bidirectional,
+        self.dropout
+    )
+
+    def bpro(x, hx, w, seq_length, out, dout):
+        y, hy, reverse, _ = out
+        dy, dhy, _, _ = dout
+        dx, dhx, dw = gru_grad_v2(x, hx, w, seq_length, y, hy, dy, dhy, reverse)
+        return dx, dhx, dw, (0)
+
+    return bpro
+
+
 @bprop_getters.register(rl_ops.CudnnGRU)
 def get_bprop_gru(self):
     """Grad definition for `GRU` operation."""
@@ -1321,7 +1226,7 @@ def get_bprop_pad(self):
         for item in paddings:
             begin += (item[0],)
         shp = shape_op(x)
-        if is_shape_unknown(shp):
+        if F.is_sequence_value_unknown(shp):
             shp = dyn_shape_op(x)
         dx = P.Slice()(dout, begin, shp)
         return (dx,)
@@ -1353,7 +1258,7 @@ def get_bprop_roi_align(self):
 
     def bprop(inputs, rois, out, dout):
         inputs_shape = shape_op(inputs)
-        if is_shape_unknown(inputs_shape):
+        if F.is_sequence_value_unknown(inputs_shape):
             inputs_shape = dyn_shape(inputs)
         dx = G.ROIAlignGrad(pooled_height, pooled_width, spatial_scale, sample_num)(dout, rois, inputs_shape)
         return dx, zeros_like(rois)
@@ -1380,7 +1285,7 @@ def get_bprop_conv2d_backprop_input(self):
 
     def bprop(x, w, f_sizes, out, dout):
         w_shape = get_shape(w)
-        if is_shape_unknown(w_shape):
+        if F.is_sequence_value_unknown(w_shape):
             w_shape = get_dyn_shape(w)
         dx = input_grad(dout, w)
         dw = filter_grad(x, dout, w_shape)
@@ -1427,8 +1332,8 @@ def get_bprop_bce_with_logits_loss(self):
             dx = mul(dx, weight)
             grad_target = mul(grad_target, weight)
         if reduction == 'mean':
-            dx_size = dyn_size(dx) if is_shape_unknown(shape(dx)) else size(dx)
-            target_size = dyn_size(target) if is_shape_unknown(shape(target)) else size(target)
+            dx_size = dyn_size(dx) if F.is_sequence_value_unknown(shape(dx)) else size(dx)
+            target_size = dyn_size(target) if F.is_sequence_value_unknown(shape(target)) else size(target)
             dx = dx / dx_size
             grad_target = grad_target / target_size
         return dx, grad_target, zeros_like(weight), zeros_like(pos_weight)
@@ -1439,18 +1344,19 @@ def get_bprop_bce_with_logits_loss(self):
 @bprop_getters.register(P.KLDivLoss)
 def get_bprop_kl_div_loss(self):
     """Grad definition for `KLDivLoss` operation."""
-    if self.reduction == "mean":
-        grad = G.KLDivLossGrad("sum")
-    else:
-        grad = G.KLDivLossGrad(self.reduction)
+    reduce_type = self.reduction
+
     size = P.Size()
     shape = P.Shape()
-    reduce_type = self.reduction
 
     def bprop(x, y, out, dout):
+        if reduce_type == "mean":
+            grad = G.KLDivLossGrad("sum")
+        else:
+            grad = G.KLDivLossGrad(self.reduction)
         dx = grad(dout, x, y)
         if reduce_type == "mean":
-            x_size = dyn_size(x) if is_shape_unknown(shape(x)) else size(x)
+            x_size = dyn_size(x) if F.is_sequence_value_unknown(shape(x)) else size(x)
             return dx / x_size, zeros_like(y)
         return dx, zeros_like(y)
 
@@ -1583,7 +1489,7 @@ def get_bprop_conv2d_backprop_filter(self):
 
     def bprop(dy, x, filter_size, out, dout):
         x_shape = get_shape(x)
-        if is_shape_unknown(x_shape):
+        if F.is_sequence_value_unknown(x_shape):
             x_shape = get_dyn_shape(x)
         dw_dx = input_grad(dy, dout, x_shape)
         dw_dy = filter_grad(x, dout)