DDIM Timestep Subset

Problem

Pick the timestep subset used by DDIM sampling.

Sample sample_steps timesteps evenly from 0 to total_steps - 1 inclusive with round(i * (total_steps - 1) / (sample_steps - 1)). Return the integer timesteps in increasing order; when sample_steps <= 1 return [0].

Test cases

case 1

in: [1000,4]

expected: [0,333,666,999]

case 2

in: [10,2]

expected: [0,9]

case 3

in: [3,1]

expected: [0]

case 4

in: [100,5]

expected: [0,25,50,74,99]

Starter codePython

def ddim_timesteps(total_steps, sample_steps):
    # Your code here
    pass

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