Sinusoidal Time Embedding

Problem

Build a DDPM-style sinusoidal timestep embedding of dimension dim (assumed even). Frequencies are exp(-log(10000) * i / (dim/2)) for i = 0 .. dim/2 - 1; the first half of the output is sin(t * freq_i) and the second half is cos(t * freq_i).

The signature is time_embedding(timestep, dim, max_period=10000.0).

Hint

Unlike transformer positional encodings, the sines and cosines are concatenated in halves here.

Test cases

case 1

in: [0,4]

expected: [0,0,1,1]

case 2

in: [1,4]

expected: [0.8414709848,0.0099998333,0.5403023059,0.9999500004]

case 3

in: [10,2]

expected: [-0.5440211109,-0.8390715291]

case 4

in: [100,6]

expected: [-0.5063656411,-0.9974947164,0.2137806661,0.8623188723,-0.0707410121,0.9768816852]

Starter codePython

import math
def time_embedding(timestep, dim, max_period=10000.0):
    # Your code here
    pass

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