b66f552
1
2
3
4
5
6
7
8
9
10
import torch @torch.jit.script def normalize_output(q: torch.Tensor, k: torch.Tensor, o: torch.Tensor) -> torch.Tensor: k = k.cumsum(1) z = (q * k).sum(-1, keepdim=True) return o / (z + 1e-10)
