pytorch · codexetreme · Apr 7, 2019 · Apr 7, 2019 · Apr 7, 2019 · Apr 9, 2019
diff --git a/aten/src/ATen/native/Unique.cpp b/aten/src/ATen/native/Unique.cpp
@@ -154,6 +154,27 @@ std::tuple<Tensor, Tensor, Tensor> _unique_dim_cpu_template(
     const bool consecutive,
     const bool return_inverse,
     const bool return_counts) {
+
+    auto sizes = self.sizes().vec();
+    // check how many zero dimensions exist
+    auto num_zero_dims = std::count(sizes.begin(), sizes.end(), 0);
+
+    // tensor is not well formed as it has 0 sized dimensions
+    if (self.size(dim) == 0){
+      AT_CHECK(
+          num_zero_dims == 1,
+          "Number of zero sized dimensions is more than one, so unique cannot be applied ")
+      Tensor output = at::empty({0}, self.options());
+      Tensor inverse_indices =
+          at::empty({0}, self.options().dtype(kLong));
+      Tensor counts = at::empty({0}, self.options().dtype(kLong));
+
+      return std::make_tuple(output, inverse_indices, counts);
+    }
+
+    AT_CHECK(num_zero_dims == 0,
+    "There are 0 sized dimensions, and they aren't selected, so unique cannot be applied");
+
   // reshape tensor as [dim, -1]
   Tensor input_flat = self.transpose(dim, 0);
   auto orig_sizes = input_flat.sizes().vec();

diff --git a/aten/src/ATen/native/cuda/Unique.cu b/aten/src/ATen/native/cuda/Unique.cu
@@ -138,6 +138,26 @@ std::tuple<Tensor, Tensor, Tensor> unique_dim_cuda_template(
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   auto allocator = THCThrustAllocator(globalContext().lazyInitCUDA());
   auto policy = thrust::cuda::par(allocator).on(stream);
+
+  auto sizes = self.sizes().vec();
+  // check how many zero dimensions exist
+  auto num_zero_dims = std::count(sizes.begin(), sizes.end(), 0);
+
+  // tensor is not well formed as it has 0 sized dimensions
+  if (self.size(dim) == 0){
+    AT_CHECK(
+        num_zero_dims == 1,
+        "Number of zero sized dimensions is more than one, so unique cannot be applied ")
+    Tensor output = at::empty({0}, self.options());
+    Tensor inverse_indices =
+        at::empty({0}, self.options().dtype(kLong));
+    Tensor counts = at::empty({0}, self.options().dtype(kLong));
+
+    return std::make_tuple(output, inverse_indices, counts);
+  }
+
+  AT_CHECK(num_zero_dims == 0,
+    "There are 0 sized dimensions, and they aren't selected, so unique cannot be applied");
 
   int64_t num_inp = self.size(dim);
   auto options = self.options().dtype(kLong);

diff --git a/test/test_torch.py b/test/test_torch.py
@@ -10848,6 +10848,9 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
                                [0., 1.]]],
                              dtype=dtype,
                              device=device)
+            x_empty = torch.empty(5, 0, dtype=dtype, device=device)
+            x_ill_formed_empty = torch.empty(5, 0, 0, dtype=dtype, device=device)
+            x_ill_formed_empty_another = torch.empty(5, 0, 5, dtype=dtype, device=device)
             expected_unique_dim0 = torch.tensor([[[1., 1.],
                                                   [0., 1.],
                                                   [2., 1.],
@@ -10878,7 +10881,9 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
                                                 device=device)
             expected_inverse_dim2 = torch.tensor([0, 1])
             expected_counts_dim2 = torch.tensor([1, 1])
-
+            expected_unique_empty = torch.tensor([], dtype=dtype, device=device)
+            expected_inverse_empty = torch.tensor([], dtype=torch.long, device=device)
+            expected_counts_empty = torch.tensor([], dtype=torch.long, device=device)
             # dim0
             x_unique = torch.unique(x, dim=0)
             self.assertEqual(expected_unique_dim0, x_unique)
@@ -10963,6 +10968,33 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
             self.assertEqual(expected_inverse_dim2, x_inverse)
             self.assertEqual(expected_counts_dim2, x_counts)
 
+            # test empty tensor
+            x_unique, x_inverse, x_counts = torch.unique(
+                x_empty,
+                return_inverse=True,
+                return_counts=True,
+                dim=1)
+            self.assertEqual(expected_unique_empty, x_unique)
+            self.assertEqual(expected_inverse_empty, x_inverse)
+            self.assertEqual(expected_counts_empty, x_counts)
+
+            # test not a well formed tensor
+            # Checking for runtime error, as this is the expected behaviour
+            with self.assertRaises(RuntimeError):
+                torch.unique(
+                    x_ill_formed_empty, 
+                    return_inverse=True, 
+                    return_counts=True, 
+                    dim=1) 
+
+            # test along dim2
+            with self.assertRaises(RuntimeError):
+                torch.unique(
+                    x_ill_formed_empty_another, 
+                    return_inverse=True, 
+                    return_counts=True, 
+                    dim=2)
+
             # test consecutive version
             y = torch.tensor(
                 [[0, 1],