Function get_timestep_embedding

diffpack/layer.py:78–85

(embedding_type, embedding_dim, embedding_scale=10000)

def get_timestep_embedding(embedding_type, embedding_dim, embedding_scale=10000):
    if embedding_type == 'sinusoidal':
        emb_func = SinusoidalEmbedding(embedding_dim, scale=embedding_scale)
    elif embedding_type == 'fourier':
        emb_func = GaussianFourierEmbedding(embedding_size=embedding_dim, scale=embedding_scale)
    else:
        raise NotImplementedError
    return emb_func