| 180 | return (sin, cos) |
| 181 | |
| 182 | def _init_weights(self): |
| 183 | device = self.periods.device |
| 184 | dtype = self.dtype |
| 185 | if self.base is not None: |
| 186 | periods = self.base ** ( |
| 187 | 2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2) |
| 188 | ) |
| 189 | else: |
| 190 | base = self.max_period / self.min_period |
| 191 | exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype) |
| 192 | periods = base**exponents |
| 193 | periods = periods / base |
| 194 | periods = periods * self.max_period |
| 195 | self.periods.data = periods |
| 196 | |
| 197 | |
| 198 | def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): |