| 518 | |
| 519 | @register_lower_rule(mops.LSQ) |
| 520 | def lsq_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 521 | assert ( |
| 522 | len(args) == 4 and len(ctx.vars_in) == 4 |
| 523 | ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}" |
| 524 | assert args[1].ndim == args[2].ndim, f"{args[1].shape}, {args[2].shape}" |
| 525 | qmax = np.array(ctx.param["qmax"], dtype=args[0].dtype) |
| 526 | qmin = np.array(ctx.param["qmin"], dtype=args[0].dtype) |
| 527 | inp, scale, zero_point, grad_s = args[:4] |
| 528 | res = inp / scale + zero_point |
| 529 | res = maximum(minimum(res, qmax), qmin) |
| 530 | res = round(res) |
| 531 | res = (res - zero_point) * scale |
| 532 | return res |
| 533 | |
| 534 | |
| 535 | @register_lower_rule("LSQBackward") |