Timestep Sinusoidal Embedding

Problem

Build the standard sinusoidal timestep embedding used by diffusion models.

For i = 0 .. dim/2 - 1 the frequency is t / 10000^(2i / dim); the first half of the output holds its sine and the second half its cosine. Assume dim is even and return the embedding rounded to 6 decimal places.

Test cases

case 1

in: [0,4]

expected: [0,0,1,1]

case 2

in: [1,4]

expected: [0.841471,0.01,0.540302,0.99995]

case 3

in: [2,2]

expected: [0.909297,-0.416147]

case 4

in: [1000,8]

expected: [0.82688,-0.506366,-0.544021,0.841471,0.562379,0.862319,-0.839072,0.540302]

Starter codePython

def timestep_embedding(t, dim):
    # Your code here
    pass

4 test cases · run your code against them in the editor