Function timestep_embedding

python/test.py:31–35

(timesteps, dim, max_period=10000)

print(out)'''

def timestep_embedding(timesteps, dim, max_period=10000):
  half = dim // 2
  freqs = (-math.log(max_period) * torch.arange(half) / half).exp()
  args = timesteps * freqs
  return torch.cat( (args.cos(), args.sin()) ).reshape(1, -1)

timesteps = Tensor([1, 2, 3]).reshape((3, 1))
dim = 10