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

kshitij12345 · facebook-github-bot · commit 0394c5a2835b · 2020-08-31T11:57:42.000-07:00
Summary: Fixes #43768 TO-DO: * [x] Add test Pull Request resolved: #43775 Reviewed By: ZolotukhinM Differential Revision: D23421979 Pulled By: ngimel fbshipit-source-id: 949fcdd30f18d17ae1c372fa6ca6a0b8d0d538ce
diff --git a/aten/src/ATen/native/Distributions.cpp b/aten/src/ATen/native/Distributions.cpp
@@ -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});
   }
diff --git a/aten/src/ATen/native/cuda/MultinomialKernel.cu b/aten/src/ATen/native/cuda/MultinomialKernel.cu
@@ -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});
   }
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -17942,13 +17942,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 == ():

Original file line number	Diff line number	Diff line change
`@@ -454,6 +454,7 @@ Tensor& multinomial_out(Tensor& result, const Tensor& self, int64_t n_sample, bo`
`454`	`454`	`if (self.dim() > 1) {`
`455`	`455`	`int64_t n_dist = self.size(-2);`
`456`	`456`	`result.resize_({n_dist, n_sample});`
	`457`	`+ if (n_dist == 0) { return result; };`
`457`	`458`	`} else {`
`458`	`459`	`result.resize_({n_sample});`
`459`	`460`	`}`
Original file line number	Diff line number	Diff line change
`@@ -395,6 +395,8 @@ void multinomial_kernel_impl(Tensor& result, const Tensor& self, const int64_t n`
`395`	`395`	`}`
`396`	`396`	`});`
`397`	`397`
	`398`	`+ AT_CUDA_CHECK(cudaGetLastError());`
	`399`	`+`
`398`	`400`	`if (inputSize == 1) {`
`399`	`401`	`result.resize_({n_sample});`
`400`	`402`	`}`