Function freq_weight_nd

k_diffusion/layers.py:36–39

(shape, scales=0, dtype=None, device=None)

@lru_cache
def freq_weight_nd(shape, scales=0, dtype=None, device=None):
    indexers = [[slice(None) if i == j else None for j in range(len(shape))] for i in range(len(shape))]
    weights = [freq_weight_1d(n, scales, dtype, device)[ix] for n, ix in zip(shape, indexers)]
    return reduce(torch.minimum, weights)


# Karras et al. preconditioned denoiser