| 117 | |
| 118 | |
| 119 | class HadamardFeatureMap(nn.Module): |
| 120 | def __init__( |
| 121 | self, |
| 122 | head_dim: int |
| 123 | ) -> HadamardFeatureMap: |
| 124 | super().__init__() |
| 125 | # Trainable map |
| 126 | self.layer1 = nn.Linear(head_dim, head_dim) |
| 127 | self.layer2 = nn.Linear(head_dim, head_dim) |
| 128 | |
| 129 | def forward(self, x: torch.Tensor): |
| 130 | return self.layer1(x) * self.layer2(x) |
| 131 | |
| 132 | |
| 133 | class LearnableOuterProductFeatureMap(nn.Module): |