Function sample

utils/stepfun.py:175–218

Piecewise-Constant PDF sampling from a step function. Args: rand: random number generator (or None for `linspace` sampling). t: [..., num_bins + 1], bin endpoint coordinates (must be sorted) w_logits: [..., num_bins], logits corresponding to bin weights num_samples: int, the numbe

(rand,
           t,
           w_logits,
           num_samples,
           single_jitter=False,
           deterministic_center=False)

def sample(rand,
           t,
           w_logits,
           num_samples,
           single_jitter=False,
           deterministic_center=False):
    """Piecewise-Constant PDF sampling from a step function.

  Args:
    rand: random number generator (or None for `linspace` sampling).
    t: [..., num_bins + 1], bin endpoint coordinates (must be sorted)
    w_logits: [..., num_bins], logits corresponding to bin weights
    num_samples: int, the number of samples.
    single_jitter: bool, if True, jitter every sample along each ray by the same
      amount in the inverse CDF. Otherwise, jitter each sample independently.
    deterministic_center: bool, if False, when `rand` is None return samples that
      linspace the entire PDF. If True, skip the front and back of the linspace
      so that the centers of each PDF interval are returned.

  Returns:
    t_samples: [batch_size, num_samples].
  """
    eps = torch.finfo(t.dtype).eps
    # eps = 1e-3

    device = t.device

    # Draw uniform samples.
    if not rand:
        if deterministic_center:
            pad = 1 / (2 * num_samples)
            u = torch.linspace(pad, 1. - pad - eps, num_samples, device=device)
        else:
            u = torch.linspace(0, 1. - eps, num_samples, device=device)
        u = torch.broadcast_to(u, t.shape[:-1] + (num_samples,))
    else:
        # `u` is in [0, 1) --- it can be zero, but it can never be 1.
        u_max = eps + (1 - eps) / num_samples
        max_jitter = (1 - u_max) / (num_samples - 1) - eps
        d = 1 if single_jitter else num_samples
        u = torch.linspace(0, 1 - u_max, num_samples, device=device) + \
            torch.rand(t.shape[:-1] + (d,), device=device) * max_jitter

    return invert_cdf(u, t, w_logits)


def sample_np(rand,