(dy, x, w, x_mean, rstd, affine, group, fmt)
| 244 | |
| 245 | |
| 246 | def group_norm_grad(dy, x, w, x_mean, rstd, affine, group, fmt): |
| 247 | assert x.ndim == 4, f"group/instance norm requires 4D input, get {x.shape}" |
| 248 | if fmt == "NCHW": |
| 249 | norm_axis = 2 |
| 250 | grouped_shape = (x.shape[0], group, -1) |
| 251 | affine_axis = 1 |
| 252 | affine_shape = (1, -1, 1, 1) |
| 253 | elif fmt == "NHWC": |
| 254 | norm_axis = 1 |
| 255 | grouped_shape = (x.shape[0], -1, group) |
| 256 | affine_axis = 1 |
| 257 | affine_shape = (1, 1, 1, -1) |
| 258 | assert False, "not checked NHWC format" |
| 259 | else: |
| 260 | assert False, f"invalid format: {fmt}" |
| 261 | |
| 262 | scaled_dy = dy if w is None else dy * w.reshape(affine_shape) |
| 263 | grouped_x = x.reshape(grouped_shape) # (N, C, H, W) -> (N, Group, -1) |
| 264 | grouped_dy = scaled_dy.reshape(grouped_shape) |
| 265 | |
| 266 | dx, dw_helper = _normalize_grad(grouped_dy, grouped_x, x_mean, rstd, norm_axis) |
| 267 | dx = dx.reshape(x.shape) |
| 268 | if affine: |
| 269 | no_affine_axis = [i for i in range(dx.ndim) if i != affine_axis] |
| 270 | db = dy.sum(no_affine_axis, False) |
| 271 | dw = (dy * dw_helper.reshape(dy.shape)).sum(no_affine_axis, False) |
| 272 | return dx, dw, db |
| 273 | return dx |
| 274 | |
| 275 | |
| 276 | @register_lower_rule(mops.GroupNorm) |
no test coverage detected