(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 192 | |
| 193 | @register_lower_rule(mops.LayerNorm) |
| 194 | def layer_norm_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 195 | assert ctx.op.normalized_dim > 0 |
| 196 | if ctx.op.affine: |
| 197 | assert len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 3 |
| 198 | x, w, b = args[0], args[1], args[2] |
| 199 | else: |
| 200 | assert len(args) == 1 and len(ctx.vars_in) == 1 and len(ctx.vars_out) == 3 |
| 201 | x, w, b = args[0], None, None |
| 202 | |
| 203 | return layer_norm(x, w, b, ctx.op.affine, ctx.op.normalized_dim, ctx.op.eps) |
| 204 | |
| 205 | |
| 206 | @register_lower_rule("LayerNormBackward") |
nothing calls this directly
no test coverage detected