(self, timesteps)
| 48 | self.scale = scale |
| 49 | |
| 50 | def forward(self, timesteps): |
| 51 | timesteps *= self.scale |
| 52 | assert timesteps.ndim == 1 |
| 53 | half_dim = self.embedding_dim // 2 |
| 54 | emb = math.log(self.max_positions) / (half_dim - 1) |
| 55 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb) |
| 56 | emb = timesteps.float()[:, None] * emb[None, :] |
| 57 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) |
| 58 | if self.embedding_dim % 2 == 1: # zero pad |
| 59 | emb = F.pad(emb, (0, 1), mode='constant') |
| 60 | assert emb.shape == (timesteps.shape[0], self.embedding_dim) |
| 61 | return emb |
| 62 | |
| 63 | |
| 64 | class GaussianFourierEmbedding(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected