Function extract_

schedule/diffusionSample.py:8–11

(a, t, x_shape)

def extract_(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def extract(v, t, x_shape):