(dy, x, x_mean, rstd, axes)
| 148 | |
| 149 | |
| 150 | def _normalize_grad(dy, x, x_mean, rstd, axes): |
| 151 | axes = _normalize_reduce_axes(x.shape, axes) |
| 152 | |
| 153 | if x_mean.ndim < x.ndim or rstd.ndim < x.ndim: |
| 154 | reduced_shape = [x.shape[i] if i not in axes else 1 for i in range(x.ndim)] |
| 155 | axes_divider = np.prod([x.shape[i] for i in axes]).astype("float32") |
| 156 | assert axes_divider != 0, f"input shape of normalize error {x.shape}" |
| 157 | |
| 158 | x_mean = x_mean.reshape(reduced_shape) |
| 159 | rstd = rstd.reshape(reduced_shape) |
| 160 | |
| 161 | delta = x - x_mean |
| 162 | reducea = dy.sum(axes, True) |
| 163 | b = (dy * delta * rstd ** 2.0).sum(axes, True) |
| 164 | x1 = dy * rstd |
| 165 | x2 = -1.0 * rstd / axes_divider * reducea |
| 166 | x3 = -1.0 * rstd / axes_divider * delta * b |
| 167 | x4 = rstd / axes_divider * (delta / axes_divider * b).sum(axes, True) |
| 168 | dx = x1 + x2 + x3 + x4 |
| 169 | return dx, delta * rstd |
| 170 | |
| 171 | |
| 172 | def layer_norm(x, w, b, affine, normalized_dim, eps): |
no test coverage detected