torch.multinomial example is incorrect

[e13h]

📚 Documentation

The torch.multinomial documentation page has an example that seems out-of-date or just wrong 🤷‍♂️

>>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights
>>> torch.multinomial(weights, 2)
tensor([1, 2])
>>> torch.multinomial(weights, 4) # ERROR!
RuntimeError: invalid argument 2: invalid multinomial distribution (with replacement=False,
not enough non-negative category to sample) at ../aten/src/TH/generic/THTensorRandom.cpp:320
>>> torch.multinomial(weights, 4, replacement=True)
tensor([ 2,  1,  1,  1])

Here is a link to the example in the code

pytorch/torch/_torch_docs.py

Lines 5588 to 5597 in 8819bad

	Example::
	>>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights
	>>> torch.multinomial(weights, 2)
	tensor([1, 2])
	>>> torch.multinomial(weights, 4) # ERROR!
	RuntimeError: invalid argument 2: invalid multinomial distribution (with replacement=False,
	not enough non-negative category to sample) at ../aten/src/TH/generic/THTensorRandom.cpp:320
	>>> torch.multinomial(weights, 4, replacement=True)
	tensor([ 2, 1, 1, 1])

The example seems to contradict the note that says

The rows of input...must be non-negative, finite and have a non-zero sum.

I am fairly sure that 0 is non-negative and finite, so as long as there are other positive inputs in the input tensor (to satisfy the non-zero sum requirement), it should work.

Here is the output that I see when I run the example code in Google Colab.

>>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights
>>> torch.multinomial(weights, 2)
tensor([1, 2])
>>> torch.multinomial(weights, 4)
tensor([2, 1, 0, 3])
>>> torch.multinomial(weights, 4, replacement=True)
tensor([2, 1, 1, 2])

I think it is also worth noting that the example shows the following runtime error

invalid multinomial distribution (with replacement=False, not enough non-negative category to sample)

Which gives me the impression that you can have negative numbers in the input tensor, and it just will not sample them. This is not the case, however. Note the observed runtime error when I change the 0's in the weight tensor to -1's:

>>> weights = torch.tensor([-1, 10, 3, -1], dtype=torch.float)
>>> torch.multinomial(weights, 2)
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

And just for the sake of being thorough, this is the runtime error that shows up when you try to take too many samples:

>>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights
>>> torch.multinomial(weights, 5)
RuntimeError: cannot sample n_sample > prob_dist.size(-1) samples without replacement
>>> torch.multinomial(weights, 5, replacement=True)
tensor([1, 2, 1, 1, 2])

cc @fritzo @neerajprad @alicanb @vishwakftw @nikitaved @jlin27 @mruberry

[mruberry]

Thanks you for reporting this issue, @evanphilipsmith.

cc @ngimel

[jihwanp]

>>> torch.multinomial(weights, 4)
tensor([2, 1, 0, 3])

Hasn't this part been revised yet? Unlike documentation, neither I don't get an error message in torch 1.10.0.

[sssilvar]

Issue description (brief recap)

In fact it is not a documentation issue it is a bug - unexpected behaviour

Steps to reproduce the issue

import torch

weights = torch.tensor([0, 10, 3, 0], dtype=torch.float)
torch.multinomial(weights, 4)

What's the expected result?

RuntimeError: invalid argument 2: invalid multinomial distribution (with replacement=False,
not enough non-negative category to sample) at ../aten/src/TH/generic/THTensorRandom.cpp:320

What's the actual result?

No errors are raised

tensor([2, 1, 0, 3])

Possible cause

This seems to be an error in an old file THTensorRandom.cpp that did all the validations including checking for zero values inside the weights n_zeros.

    THArgCheckWithCleanup((with_replacement || (n_categories - n_zeros >= n_sample)),
                          THCleanup(THDoubleTensor_free(cum_dist); if (start_dim == 1) THTensor_(squeeze1d)(prob_dist, prob_dist, 0);),
                          2,
                          "invalid multinomial distribution (with replacement=False, not enough non-negative category to sample)");

Now the file in charge of checking for the right behavior is Distributions.cpp and the corresponding assertion for zero values (TORCH_CHECK) is absent.

[ptrblck]

Also reported here.

[Ignasijus]

Hey, 3 years have passed, why this bug is still not fixed?