(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 534 | |
| 535 | @register_lower_rule("LSQBackward") |
| 536 | def lsq_grad_rule(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 537 | assert ( |
| 538 | len(args) == 5 and len(ctx.vars_in) == 5 |
| 539 | ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}" |
| 540 | assert args[0].ndim == args[1].ndim, f"{args[0].shape}, {args[1].shape}" |
| 541 | |
| 542 | inp = args[1] / args[2] + args[3] |
| 543 | round_inp = round(inp) |
| 544 | qmax = np.array(ctx.param["qmax"], dtype=inp.dtype) |
| 545 | qmin = np.array(ctx.param["qmin"], dtype=inp.dtype) |
| 546 | ind_small = inp < qmin |
| 547 | ind_big = inp > qmax |
| 548 | ind_middle = logical_or(ind_small, ind_big) |
| 549 | ind_middle = logical_not(ind_middle) |
| 550 | ind_small = ind_small.astype(inp.dtype) |
| 551 | ind_big = ind_big.astype(inp.dtype) |
| 552 | |
| 553 | grad_s = ind_small * qmin + ind_big * qmax + ind_middle * (-inp + round_inp) |
| 554 | grad_s = grad_s * args[-1] * args[0] |
| 555 | grad_x = ind_middle * args[0] |
| 556 | return [grad_x, grad_s] |
| 557 | |
| 558 | |
| 559 | @register_lower_rule(mops.CheckNonFinite) |
nothing calls this directly
no test coverage detected