(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 275 | |
| 276 | @register_lower_rule(mops.GroupNorm) |
| 277 | def group_norm_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 278 | if ctx.op.affine: |
| 279 | assert len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 3 |
| 280 | x, w, b = args[0], args[1], args[2] |
| 281 | else: |
| 282 | assert len(args) == 1 and len(ctx.vars_in) == 1 and len(ctx.vars_out) == 3 |
| 283 | x, w, b = args[0], None, None |
| 284 | |
| 285 | assert x.ndim == 4, f"group norm requires 4D input, get {x.shape}" |
| 286 | |
| 287 | if ctx.op.format == mops.GroupNorm.Format.NCHW: |
| 288 | return group_norm(x, w, b, ctx.op.affine, ctx.op.eps, ctx.op.group, "NCHW") |
| 289 | elif ctx.op.format == mops.GroupNorm.Format.NHWC: |
| 290 | return group_norm(x, w, b, ctx.op.affine, ctx.op.eps, ctx.op.group, "NHWC") |
| 291 | else: |
| 292 | assert False, f"invalid format: {ctx.op.format}" |
| 293 | |
| 294 | |
| 295 | @register_lower_rule("GroupNormBackward") |
nothing calls this directly
no test coverage detected