choice

deeptrack.backend.array_api_compat_ext.torch.random.choice(a: int | Tensor, size: tuple[int, ...] | None = None, replace: bool = True, p: Tensor | None = None) → Tensor

Sample from a 1D tensor or from range(a).

This function mirrors numpy.random.choice.

Parameters

a: int | torch.Tensor

If an integer, samples are drawn from torch.arange(a). If a tensor, it must be 1D and samples are drawn from its elements.

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

Output shape. If None, returns a scalar 0D tensor.

replace: bool, optional

Whether sampling is with replacement. Defaults to True.

p: torch.Tensor | None, optional

Optional probability weights. Must have the same length as the population and sum to 1 (normalization is applied internally).

Returns

torch.Tensor

Samples drawn from a (or from range(a) if a is an integer).

Raises

ValueError

If a is a tensor and is not 1D, if a is an integer < 1, or if p has an incompatible shape.

Examples

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

Sample a scalar from a tensor:

>>> import torch
>>>
>>> a = torch.tensor([10, 20, 30, 40])
>>> rnd.choice(a)
tensor(40)

Sample an array of shape (2, 3):

>>> rnd.choice(a, (2, 3)).shape
torch.Size([2, 3])

Sample from range(5) (NumPy parity with np.random.choice(5)):

>>> rnd.choice(5, (4,)).shape
torch.Size([4])

Use probabilities (always pick index 2 from range(4)):

>>> p = torch.tensor([0.0, 0.0, 1.0, 0.0])
>>> rnd.choice(4, (3,), p=p)
tensor([2, 2, 2])