Function extract_into_tensor

diff2flow/utils/diffusion_utils.py:56–59

(a, t, x_shape)

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

def extract_and_interpolate_into_tensor(a, t, x_shape):
    b, *_ = t.shape