| 165 | """ |
| 166 | |
| 167 | def __init__( |
| 168 | self, |
| 169 | E_q: int, |
| 170 | E_k: int, |
| 171 | E_v: int, |
| 172 | E_total: int, |
| 173 | nheads: int, |
| 174 | dropout: float = 0.0, |
| 175 | bias=True, |
| 176 | device=None, |
| 177 | dtype=None, |
| 178 | batch_first=False, |
| 179 | ): |
| 180 | factory_kwargs = {"device": device, "dtype": dtype} |
| 181 | super().__init__() |
| 182 | self.nheads = nheads |
| 183 | self.dropout = dropout |
| 184 | self._qkv_same_embed_dim = E_q == E_k and E_q == E_v |
| 185 | if self._qkv_same_embed_dim: |
| 186 | self.packed_proj = nn.Linear(E_q, E_total * 3, bias=bias, **factory_kwargs) |
| 187 | else: |
| 188 | self.q_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) |
| 189 | self.k_proj = nn.Linear(E_k, E_total, bias=bias, **factory_kwargs) |
| 190 | self.v_proj = nn.Linear(E_v, E_total, bias=bias, **factory_kwargs) |
| 191 | E_out = E_q |
| 192 | self.out_proj = nn.Linear(E_total, E_out, bias=bias, **factory_kwargs) |
| 193 | assert E_total % nheads == 0, "Embedding dim is not divisible by nheads" |
| 194 | self.E_head = E_total // nheads |
| 195 | self.bias = bias |
| 196 | self.batch_first = batch_first |
| 197 | |
| 198 | def forward( |
| 199 | self, |