Introduction

PyTorch does not have a single drop-in function that behaves exactly like numpy.random.choice in every case, but it does provide the building blocks you usually need. The right tool depends on whether you want to sample indices uniformly, sample values from a tensor, shuffle without replacement, or perform weighted random sampling.

In practice, the most useful APIs are torch.randint, torch.randperm, and torch.multinomial. Once you understand the difference between them, “random choice” in PyTorch becomes straightforward.

Use torch.randint for Uniform Sampling with Replacement

If you want random indices where the same index can appear more than once, torch.randint is the simplest option.

1import torch
2
3torch.manual_seed(42)
4
5values = torch.tensor([10, 20, 30, 40, 50])
6indices = torch.randint(low=0, high=len(values), size=(4,))
7sample = values[indices]
8
9print(indices)
10print(sample)

tensor([2, 2, 1, 4])
tensor([30, 30, 20, 50])

This is uniform sampling with replacement. Each draw is independent, so duplicates are allowed.

Use torch.randperm to Shuffle Without Replacement

When you want a random ordering of all positions, use torch.randperm.

1import torch
2
3torch.manual_seed(7)
4
5values = torch.tensor([10, 20, 30, 40, 50])
6order = torch.randperm(len(values))
7shuffled = values[order]
8
9print(order)
10print(shuffled)

If you only need k unique choices, slice the permutation:

1import torch
2
3values = torch.tensor([10, 20, 30, 40, 50])
4k = 3
5indices = torch.randperm(len(values))[:k]
6sample = values[indices]
7
8print(sample)

This is the usual answer when you need random unique elements.

Use torch.multinomial for Choice-Like Sampling

torch.multinomial is the closest PyTorch equivalent to “choose items according to probabilities.” It works with a tensor of non-negative weights.

1import torch
2
3torch.manual_seed(0)
4
5values = torch.tensor(["red", "green", "blue"])
6weights = torch.tensor([0.1, 0.2, 0.7])
7
8indices = torch.multinomial(weights, num_samples=5, replacement=True)
9print(indices)

The code above shows the sampling step, but PyTorch tensors cannot store Python strings directly in this form. In real code, keep labels in a Python list and sample indices:

1import torch
2
3torch.manual_seed(0)
4
5labels = ["red", "green", "blue"]
6weights = torch.tensor([0.1, 0.2, 0.7], dtype=torch.float32)
7
8indices = torch.multinomial(weights, num_samples=5, replacement=True)
9chosen = [labels[i] for i in indices.tolist()]
10
11print(chosen)

If replacement=False, each index can appear at most once, assuming you do not request more samples than there are positive-weight entries.

Sample Rows from a Tensor

A very common use case is randomly selecting rows from a dataset tensor.

1import torch
2
3torch.manual_seed(123)
4
5data = torch.tensor(
6    [
7        [1.0, 1.1],
8        [2.0, 2.2],
9        [3.0, 3.3],
10        [4.0, 4.4],
11    ]
12)
13
14row_indices = torch.randperm(data.size(0))[:2]
15batch = data[row_indices]
16
17print(row_indices)
18print(batch)

This pattern is useful for quick experiments, small custom training loops, or subsampling data for visualization.

Reproducibility Matters

If you want stable results for debugging, set a seed with torch.manual_seed.

1import torch
2
3torch.manual_seed(99)
4print(torch.randint(0, 10, (5,)))

Without a seed, repeated runs will generate different samples, which is usually what you want in training but not always what you want in tests.

Which Function Should You Use?

Use torch.randint when repeated picks are fine and you want raw random indices. Use torch.randperm when you need a shuffled order or unique samples. Use torch.multinomial when some choices should be more likely than others.

That mapping is more useful than searching for one magical “PyTorch choice” function, because the correct semantics matter more than the name.

Common Pitfalls

One common mistake is expecting torch.randperm to support weighted sampling. It does not; it simply returns a uniform random permutation.

Another issue is forgetting the difference between sampling with and without replacement. If duplicates appear unexpectedly, check whether you used replacement=True.

People also sometimes try to store strings in a numeric tensor and then wonder why the example fails. PyTorch tensors are numeric containers, so sampling labels usually means sampling indices first and then mapping those indices back to Python objects or another structure.

Finally, if your weights tensor contains negative values or all zeros, torch.multinomial will fail. Validate the weights before sampling.

Summary

• PyTorch random choice is usually built from torch.randint, torch.randperm, or torch.multinomial.
• Use torch.randint for uniform sampling with replacement.
• Use torch.randperm for shuffling or unique sampling without replacement.
• Use torch.multinomial for weighted sampling.
• Set torch.manual_seed when you need reproducible random results.