Function stratified_with_settings

k_diffusion/utils.py:314–321

Draws stratified samples from a uniform distribution, using settings from a context manager.

(shape, dtype=None, device=None)

def stratified_with_settings(shape, dtype=None, device=None):
    """Draws stratified samples from a uniform distribution, using settings from a context
    manager."""
    if not hasattr(stratified_settings, 'disable') or stratified_settings.disable:
        return torch.rand(shape, dtype=dtype, device=device)
    return stratified_uniform(
        shape, stratified_settings.group, stratified_settings.groups, dtype=dtype, device=device
    )


def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):