(q, k, scale, eps)
| 109 | |
| 110 | @flags.compile_wrap |
| 111 | def scale_for_cosine_sim(q, k, scale, eps): |
| 112 | dtype = reduce(torch.promote_types, (q.dtype, k.dtype, scale.dtype, torch.float32)) |
| 113 | sum_sq_q = torch.sum(q.to(dtype)**2, dim=-1, keepdim=True) |
| 114 | sum_sq_k = torch.sum(k.to(dtype)**2, dim=-1, keepdim=True) |
| 115 | sqrt_scale = torch.sqrt(scale.to(dtype)) |
| 116 | scale_q = sqrt_scale * torch.rsqrt(sum_sq_q + eps) |
| 117 | scale_k = sqrt_scale * torch.rsqrt(sum_sq_k + eps) |
| 118 | return q * scale_q.to(q.dtype), k * scale_k.to(k.dtype) |
| 119 | |
| 120 | |
| 121 | @flags.compile_wrap |
no outgoing calls
no test coverage detected