MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / fakequant_grad_rule

Function fakequant_grad_rule

imperative/python/megengine/xla/rules/math.py:464–477  ·  view source on GitHub ↗
(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])

Source from the content-addressed store, hash-verified

462
463@register_lower_rule("FakeQuantBackward")
464def fakequant_grad_rule(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
465 assert (
466 len(args) == 4 and len(ctx.vars_in) == 4
467 ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}"
468 assert args[0].ndim == args[1].ndim, f"{args[0].shape}, {args[1].shape}"
469 assert args[2].ndim == args[3].ndim, f"{args[2].shape}, {args[3].shape}"
470 diff, inp, scale, zere_point = args[:4]
471 x = round(inp / scale) + zere_point
472 qmax = np.array(ctx.param["qmax"], dtype=x.dtype)
473 qmin = np.array(ctx.param["qmin"], dtype=x.dtype)
474 mask1 = x <= qmax
475 mask2 = x >= qmin
476 mask = logical_and(mask1, mask2).astype(diff.dtype)
477 return diff * mask
478
479
480@register_lower_rule(mops.TQT)

Callers

nothing calls this directly

Calls 4

logical_andFunction · 0.85
arrayMethod · 0.80
roundFunction · 0.50
astypeMethod · 0.45

Tested by

no test coverage detected