| 479 | |
| 480 | @register_lower_rule(mops.TQT) |
| 481 | def tqt_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 482 | assert ( |
| 483 | len(args) == 2 and len(ctx.vars_in) == 2 |
| 484 | ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}" |
| 485 | qmax = np.array(ctx.param["qmax"], dtype=args[0].dtype) |
| 486 | qmin = np.array(ctx.param["qmin"], dtype=args[0].dtype) |
| 487 | t = pow(np.array(2, dtype=args[1].dtype), args[1]) |
| 488 | res = round(args[0] / t) |
| 489 | res = maximum(minimum(res, qmax), qmin) |
| 490 | return res * t |
| 491 | |
| 492 | |
| 493 | @register_lower_rule("TQTBackward") |