Computes multi-head attention. Supports nested or padded tensors. Args: E_q (int): Size of embedding dim for query E_k (int): Size of embedding dim for key E_v (int): Size of embedding dim for value E_total (int): Total embedding dim of combined heads post i
| 150 | raise NotImplementedError |
| 151 | |
| 152 | class MultiHeadAttention(nn.Module): |
| 153 | """ |
| 154 | Computes multi-head attention. Supports nested or padded tensors. |
| 155 | |
| 156 | Args: |
| 157 | E_q (int): Size of embedding dim for query |
| 158 | E_k (int): Size of embedding dim for key |
| 159 | E_v (int): Size of embedding dim for value |
| 160 | E_total (int): Total embedding dim of combined heads post input projection. Each head |
| 161 | has dim E_total // nheads |
| 162 | nheads (int): Number of heads |
| 163 | dropout (float, optional): Dropout probability. Default: 0.0 |
| 164 | bias (bool, optional): Whether to add bias to input projection. Default: True |
| 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, |
| 200 | query: torch.Tensor, |
| 201 | key: torch.Tensor, |
| 202 | value: torch.Tensor, |
| 203 | attn_mask=None, |
| 204 | is_causal=False, |
| 205 | need_weights=False, # for compatibility with nn.MultiheadAttention |
| 206 | ) -> torch.Tensor: |
| 207 | """ |
| 208 | Args: |
| 209 | query (torch.Tensor): query of shape (``N``, ``L_q``, ``E_qk``) |