zsjTiger
/

zsjTiger's picture
Duplicate from silx-ai/Quasar-Preview
5ebdfe3
Raw
History Blame Contribute Delete
190 Bytes
import torch
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)