(q, k, scale, eps)
| 37 | |
| 38 | |
| 39 | def scale_for_cosine_sim(q, k, scale, eps): |
| 40 | dtype = reduce(torch.promote_types, (q.dtype, k.dtype, scale.dtype, torch.float32)) |
| 41 | sum_sq_q = torch.sum(q.to(dtype)**2, dim=-1, keepdim=True) |
| 42 | sum_sq_k = torch.sum(k.to(dtype)**2, dim=-1, keepdim=True) |
| 43 | sqrt_scale = torch.sqrt(scale.to(dtype)) |
| 44 | scale_q = sqrt_scale * torch.rsqrt(sum_sq_q + eps) |
| 45 | scale_k = sqrt_scale * torch.rsqrt(sum_sq_k + eps) |
| 46 | return q * scale_q.to(q.dtype), k * scale_k.to(k.dtype) |
| 47 | |
| 48 | def scale_for_cosine_sim_qkv(qkv, scale, eps): |
| 49 | q, k, v = qkv.unbind(2) |
no outgoing calls
no test coverage detected