Sampling Functionals#

Weighted Multinomial#

physicsnemo.nn.functional.weighted_multinomial(
input: int | Tensor,
num_samples: int,
replacement: bool = False,
*,
generator: Generator | None = None,
strategy: Literal['exact', 'poisson_gap'] = 'exact',
device: device | str | None = None,
implementation: Literal['torch'] | None = None,
) → Tensor#

Sample indices from a weighted or uniform multinomial distribution.

The core call follows torch.multinomial(): pass input, num_samples, replacement, and optionally generator. A one-dimensional tensor input contains relative sampling weights. Unlike torch.multinomial(), input may also be an integer population size, which represents uniform weights without materializing a tensor. Batched weight tensors and the out argument are not supported.

Without replacement, the default "exact" strategy produces exact weighted or uniform samples. Moderate uniform populations use torch.randperm(); weighted and very large uniform populations use a chunked exponential race without the torch.multinomial() \(2^{24}\) category limit. Chunking avoids a temporary allocation proportional to the full population when only a small sample is requested.

The opt-in "poisson_gap" strategy draws and normalizes exponential gaps. It uses \(O(\text{num_samples})\) memory regardless of population size, but returns ordered, near-uniform coverage rather than an exact uniform random subset. It is only available for integer input without replacement.

With replacement, integer input uses torch.randint() and tensor input uses torch.multinomial().

Parameters:

  • input (int or torch.Tensor) – An integer uniform population size or one-dimensional floating-point relative weights.

  • num_samples (int) – Number of indices to sample.

  • replacement (bool, default=False) – Whether an index may be sampled more than once.

  • generator (torch.Generator, optional) – Generator used for random draws. Its device must match the sampling device.

  • strategy ({"exact", "poisson_gap"}, default="exact") – Sampling strategy. Approximate Poisson-gap sampling must be requested explicitly and is only valid for uniform sampling without replacement.

  • device (torch.device or str, optional) – Output device when input is an integer. When weights are supplied, their device is used and an explicitly supplied device must match it.

  • implementation ({"torch"}, optional) – Backend implementation. None selects the default implementation.

Returns:

Sampled indices with shape (num_samples,) and dtype torch.int64. Exact samples are returned in random draw order; Poisson-gap samples are returned in increasing index order.

Return type:

torch.Tensor

The core API follows torch.multinomial while also accepting an integer uniform population size. Without replacement, the default exact strategy provides uniform or weighted sampling without the \(2^{24}\) category limit imposed by torch.multinomial. The poisson_gap strategy is an explicit, low-memory approximation for unweighted, ordered coverage of very large populations. Sampling with replacement is also supported.