(self,
offset: Union[int, torch.Tensor],
size: int,
apply_dropout: bool = True)
| 229 | return self.dropout(x), pos_emb |
| 230 | |
| 231 | def position_encoding(self, |
| 232 | offset: Union[int, torch.Tensor], |
| 233 | size: int, |
| 234 | apply_dropout: bool = True) -> torch.Tensor: |
| 235 | |
| 236 | pe = torch.view_as_complex(self.pe) |
| 237 | if isinstance(offset, int): |
| 238 | assert offset + size <= self.max_len |
| 239 | pos_emb = pe[:, offset:offset + size] |
| 240 | else: |
| 241 | assert torch.max(offset) + size <= self.max_len |
| 242 | index = offset.unsqueeze(1) + torch.arange(0, size).to( |
| 243 | offset.device) # B X T |
| 244 | flag = index > 0 |
| 245 | # remove negative offset |
| 246 | index = index * flag |
| 247 | pos_emb = F.embedding(index, pe[0]) # B X T X head_dim//2 |
| 248 | if apply_dropout: |
| 249 | # NOTE(Mddct) dropout don't suuport complex float for pos_emb |
| 250 | pos_emb = self.dropout_complex(pos_emb) |
| 251 | return pos_emb |
| 252 | |
| 253 | def dropout_complex(self, x): |
| 254 | mask = torch.nn.functional.dropout( |
no test coverage detected