Function scale_for_cosine_sim_qkv

k_diffusion/models/modules.py:122–125

(qkv, scale, eps)

@flags.compile_wrap
def scale_for_cosine_sim_qkv(qkv, scale, eps):
    q, k, v = qkv.unbind(2)
    q, k = scale_for_cosine_sim(q, k, scale[:, None], eps)
    return torch.stack((q, k, v), dim=2)


class ConvGEGLU(nn.Module):