Function ref_grad_func

tests/test_cuda_kernel.py:67–78

(value: torch.Tensor, dy: torch.Tensor, 
                      scale: torch.Tensor, offset: torch.Tensor)

def __TEST_QUANTIZE_LT_B__(size: List[int], iterations: int, sym: bool):
    def ref_grad_func(value: torch.Tensor, dy: torch.Tensor, 
                      scale: torch.Tensor, offset: torch.Tensor):
        
        qt = ppq_tensor_round(value / scale, policy=ROUNDING_POLICY) + offset
        clipped_qt = qt.clip(Q_MIN, Q_MAX)
        
        dx = torch.where(clipped_qt != qt, torch.zeros_like(dy), dy)
        ds = torch.where(clipped_qt == qt, (((qt - offset) * scale) - value) * dy / scale, torch.zeros_like(dy))
        ds += torch.where(qt > Q_MAX, (Q_MAX - offset) * dy, torch.zeros_like(dy))
        ds += torch.where(qt < Q_MIN, (Q_MIN - offset) * dy, torch.zeros_like(dy))
        ds = ds.sum() / sqrt(value.numel() * (Q_MAX - Q_MIN))
        return dx, ds

    for _ in tqdm(range(iterations), desc='QUANTIZE LT_B TESTING...'):
        t = torch.rand(size=size).cuda() * 50