Function diff_step

custom_samplers.py:219–224

(a, b, s)

@torch.no_grad()
def diff_step(a, b, s):
    n = torch.linalg.matrix_norm(a, keepdim=True)
    x = a.div(n)
    y = b.div(torch.linalg.matrix_norm(b, keepdim=True))
    y = n * y.sub(x.mul(torch.mul(x, y).sum().clamp(min=-1.0, max=1.0)))
    return a - y * s

def perp_step_wrap(s=0.5):
    @torch.no_grad()