Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return:
(timesteps, dim, max_period=10000, repeat_only=False)
| 72 | |
| 73 | |
| 74 | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): |
| 75 | """ |
| 76 | Create sinusoidal timestep embeddings. |
| 77 | :param timesteps: a 1-D Tensor of N indices, one per batch element. |
| 78 | These may be fractional. |
| 79 | :param dim: the dimension of the output. |
| 80 | :param max_period: controls the minimum frequency of the embeddings. |
| 81 | :return: an [N x dim] Tensor of positional embeddings. |
| 82 | """ |
| 83 | if not repeat_only: |
| 84 | half = dim // 2 |
| 85 | freqs = torch.exp( |
| 86 | -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half |
| 87 | ).to(device=timesteps.device) |
| 88 | args = timesteps[:, None].float() * freqs[None] |
| 89 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| 90 | if dim % 2: |
| 91 | embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| 92 | else: |
| 93 | embedding = repeat(timesteps, 'b -> b d', d=dim).contiguous() |
| 94 | return embedding |
| 95 | |
| 96 | |
| 97 | def zero_module(module): |
nothing calls this directly
no outgoing calls
no test coverage detected