r""" Simple linear mapping feature map as in `Finetuning Pretrained Transformers into RNNs `_
| 64 | |
| 65 | |
| 66 | class T2RFeatureMap(nn.Module): |
| 67 | |
| 68 | r""" |
| 69 | Simple linear mapping feature map as in |
| 70 | `Finetuning Pretrained Transformers into RNNs <https://arxiv.org/abs/2103.13076>`_ |
| 71 | """ |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | head_dim: int, |
| 76 | dot_dim: int = None, |
| 77 | bias: Optional[bool] = False |
| 78 | ) -> T2RFeatureMap: |
| 79 | super().__init__() |
| 80 | # Trainable map |
| 81 | if dot_dim is None: |
| 82 | dot_dim = head_dim |
| 83 | |
| 84 | self.head_dim = head_dim |
| 85 | self.dot_dim = dot_dim |
| 86 | self.bias = bias |
| 87 | |
| 88 | self.layer = nn.Linear(head_dim, dot_dim, bias=bias) |
| 89 | |
| 90 | def __repr__(self) -> str: |
| 91 | return f"{self.__class__.__name__}(head_dim={self.head_dim}, dot_dim={self.dot_dim}, bias={self.bias})" |
| 92 | |
| 93 | def forward(self, x: torch.Tensor): |
| 94 | return self.layer(x).relu() |
| 95 | |
| 96 | |
| 97 | class DPFPFeatureMap(nn.Module): |