Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[fix] torch.multinomial : fix for 0 size dim #43775

Closed
wants to merge 5 commits into from
Closed

[fix] torch.multinomial : fix for 0 size dim #43775

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e07c9fd
fix for 0 size axis
kshitij12345 Aug 28, 2020
8ba7a93
Merge branch 'master' into fix/multinomial/0-dist
kshitij12345 Aug 28, 2020
0b21262
update existing test
kshitij12345 Aug 28, 2020
c2bc6d7
add AT_CUDA_CHECK
kshitij12345 Aug 29, 2020
6e1ed7d
update test
kshitij12345 Aug 29, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions aten/src/ATen/native/Distributions.cpp
View file
Open in desktop
Expand Up @@ -454,6 +454,7 @@ Tensor& multinomial_out(Tensor& result, const Tensor& self, int64_t n_sample, bo
if (self.dim() > 1) {
int64_t n_dist = self.size(-2);
result.resize_({n_dist, n_sample});
if (n_dist == 0) { return result; };
} else {
result.resize_({n_sample});
}
Expand Down
2 changes: 2 additions & 0 deletions aten/src/ATen/native/cuda/MultinomialKernel.cu
View file
Open in desktop
Expand Up @@ -395,6 +395,8 @@ void multinomial_kernel_impl(Tensor& result, const Tensor& self, const int64_t n
}
});

AT_CUDA_CHECK(cudaGetLastError());

if (inputSize == 1) {
result.resize_({n_sample});
}
Expand Down
18 changes: 12 additions & 6 deletions test/test_torch.py
View file
Open in desktop
Expand Up @@ -17796,13 +17796,19 @@ def test(probs, replacement):
test(y, True)
test(z, True)

def test_multinomial_empty(self, device):
probs = torch.ones(0, 3)
num_samples = 1
def _test_multinomial_empty(self, device, replacement, num_samples):
probs = torch.ones(0, 3, device=device)
expected = torch.empty(0, num_samples, dtype=torch.int64)
for replacement in (True, False):
out = torch.multinomial(probs, num_samples=num_samples, replacement=replacement)
self.assertEqual(out, expected)
out = torch.multinomial(probs, num_samples=num_samples, replacement=replacement)
self.assertEqual(out, expected)

def test_multinomial_empty_w_replacement(self, device):
self._test_multinomial_empty(device, True, 1)
self._test_multinomial_empty(device, True, 2)

def test_multinomial_empty_wo_replacement(self, device):
self._test_multinomial_empty(device, False, 1)
self._test_multinomial_empty(device, False, 2)

def _generate_input(self, shape, dtype, device, with_extremal):
if shape == ():
Expand Down