(dy, x, w, x_mean, rstd, affine)
| 178 | |
| 179 | |
| 180 | def layer_norm_grad(dy, x, w, x_mean, rstd, affine): |
| 181 | reduce_axes = list(range(x_mean.ndim, x.ndim)) |
| 182 | scaled_dy = dy * w if w is not None else dy |
| 183 | dx, dw_helper = _normalize_grad(scaled_dy, x, x_mean, rstd, reduce_axes) |
| 184 | |
| 185 | if affine: |
| 186 | unreduce_axes = list(range(x.ndim - w.ndim)) |
| 187 | dbias = dy.sum(unreduce_axes, keepdims=False) |
| 188 | dweight = (dw_helper * dy).sum(unreduce_axes, keepdims=False) |
| 189 | return dx, dweight, dbias |
| 190 | return dx |
| 191 | |
| 192 | |
| 193 | @register_lower_rule(mops.LayerNorm) |
no test coverage detected