Function rotate_half

k_diffusion/models/axial_rope.py:13–17

(x)

def rotate_half(x):
    x1, x2 = x[..., 0::2], x[..., 1::2]
    x = torch.stack((-x2, x1), dim=-1)
    *shape, d, r = x.shape
    return x.view(*shape, d * r)


@flags.compile_wrap