position_hint.py

class CrossPositionAttentionStem(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.scale = dim ** -0.5

    def forward(self, pe_q, pe_k, kv):
        # Compute position weighted attention
        attn = pe_q @ pe_k.transpose(-2, -1) * self.scale 
        attn = attn.softmax(dim=-1)
        return attn @ kv