multinomial

deeptrack.backend.array_api_compat_ext.torch.random.multinomial(n: int | Tensor, pvals: Tensor, size: tuple[int, ...] | None = None) → Tensor

Sample from a multinomial distribution.

Mirrors numpy.random.multinomial.

Parameters

n: int | torch.Tensor

Number of trials.

pvals: torch.Tensor

1D tensor of category probabilities.

size: tuple[int, …] | None, optional

Sample shape. If None, returns a single draw.

Returns

torch.Tensor

Counts per category. Output shape is (len(pvals),) if size=None, otherwise size + (len(pvals),). The returned dtype is torch.int64 to match NumPy parity.

Examples

>>> import deeptrack.backend.array_api_compat_ext.torch.random as rnd

Single draw:

>>> import torch
>>>
>>> p = torch.tensor([0.2, 0.8])
>>> rnd.multinomial(5, p)
tensor([1., 4.])

>>> rnd.multinomial(5, p).dtype
torch.int64

Multiple draws:

>>> rnd.multinomial(5, p, (3,)).shape
torch.Size([3, 2])