Skip to content

Commit 4123d6d

Browse files
codexetremecodexetreme
committed
Fixed review suggestions
1 parent 62a11ca commit 4123d6d

File tree

3 files changed

+56
-44
lines changed

3 files changed

+56
-44
lines changed

‎aten/src/ATen/native/Unique.cpp‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -154,28 +154,28 @@ std::tuple<Tensor, Tensor, Tensor> _unique_dim_cpu_template(
154154
const bool consecutive,
155155
const bool return_inverse,
156156
const bool return_counts) {
157-
// check if dim is in range and see if the size along that dim is 0
158-
if (dim <= self.dim() && self.size(dim) == 0) {
159-
// assume tensor has all dimentions non-zero
160-
bool well_formed_tensor = true;
161-
for (int64_t i = 0; i < self.dim(); ++i) {
162-
if (dim == i)
163-
continue;
164-
if (self.size(i) == 0) {
165-
// found atleast 1 zero sized dimention, hence the tensor is not well formed
166-
well_formed_tensor = false;
167-
break;
168-
}
169-
}
170-
if (well_formed_tensor) {
171-
Tensor output = at::empty({0}, self.options().dtype(kLong));
157+
158+
auto sizes = self.sizes().vec();
159+
// check how many zero dimentions exist
160+
auto num_zero_dims = std::count(sizes.begin(), sizes.end(),0);
161+
162+
// tensor is not well formed as it has 0 sized dimentions
163+
if (self.size(dim)==0){
164+
AT_CHECK(
165+
num_zero_dims == 1,
166+
"Number of zero sized dimentions is more than one, so unique cannot be applied ")
167+
Tensor output = at::empty({0}, self.options());
172168
Tensor inverse_indices =
173-
at::empty(self.sizes(), self.options().dtype(kLong));
174-
Tensor counts = at::empty({0}, self.options().dtype(kLong));
169+
at::empty(self.size(dim), self.options().dtype(kLong));
170+
Tensor counts = at::empty({0}, self.options().dtype(kLong));
175171

176-
return std::make_tuple(output, inverse_indices,counts);
172+
return std::make_tuple(output, inverse_indices, counts);
177173
}
178-
}
174+
if (self.size(dim) != 0 ){
175+
AT_CHECK(num_zero_dims==0,
176+
"There are 0 sized dimentions, and they aren't selected, so unique cannot be applied");
177+
}
178+
179179
// reshape tensor as [dim, -1]
180180
Tensor input_flat = self.transpose(dim, 0);
181181
auto orig_sizes = input_flat.sizes().vec();

‎aten/src/ATen/native/cuda/Unique.cu‎

Lines changed: 20 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -138,30 +138,27 @@ std::tuple<Tensor, Tensor, Tensor> unique_dim_cuda_template(
138138
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
139139
auto allocator = THCThrustAllocator(globalContext().lazyInitCUDA());
140140
auto policy = thrust::cuda::par(allocator).on(stream);
141-
if (dim <= self.dim() && self.size(dim) == 0) {
142-
bool well_formed_tensor = true;
143-
for (int64_t i = 0; i < self.dim(); ++i) {
144-
if (dim == i)
145-
continue;
146-
if (self.size(i) == 0) {
147-
// found atleast 1 zero sized dimention, hence the tensor is not well formed
148-
well_formed_tensor = false;
149-
break;
150-
}
151-
}
152-
if (well_formed_tensor) {
153-
Tensor output = at::empty({0},self.options().dtype(kLong));
154-
Tensor inverse_indices =
155-
at::empty(self.sizes(), self.options().dtype(kLong));
156-
Tensor counts = at::empty({0}, self.options().dtype(kLong));
157-
158-
THCudaCheck(cudaGetLastError());
159-
return std::make_tuple(output, inverse_indices,counts);
160-
}
141+
142+
auto sizes = self.sizes().vec();
143+
// check how many zero dimentions exist
144+
auto num_zero_dims = std::count(sizes.begin(), sizes.end(),0);
145+
146+
// tensor is not well formed as it has 0 sized dimentions
147+
if (self.size(dim)==0){
148+
AT_CHECK(
149+
num_zero_dims == 1,
150+
"Number of zero sized dimentions is more than one, so unique cannot be applied ")
151+
Tensor output = at::empty({0}, self.options());
152+
Tensor inverse_indices =
153+
at::empty(self.size(dim), self.options().dtype(kLong));
154+
Tensor counts = at::empty({0}, self.options().dtype(kLong));
155+
156+
return std::make_tuple(output, inverse_indices, counts);
157+
}
158+
if (self.size(dim) != 0 ){
159+
AT_CHECK(num_zero_dims==0,
160+
"There are 0 sized dimentions, and they aren't selected, so unique cannot be applied");
161161
}
162-
Tensor input_flat_ = self.transpose(dim, 0);
163-
auto orig_sizes = input_flat_.sizes().vec();
164-
input_flat_ = input_flat_.contiguous().view({input_flat_.size(0), -1});
165162

166163
int64_t num_inp = self.size(dim);
167164
auto options = self.options().dtype(kLong);

‎test/test_torch.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10828,6 +10828,8 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
1082810828
dtype=dtype,
1082910829
device=device)
1083010830
x_empty = torch.empty(5, 0, dtype=dtype, device=device)
10831+
x_ill_formed_empty = torch.empty(5, 0, 0, dtype=dtype, device=device)
10832+
x_ill_formed_empty_another = torch.empty(5, 0, 5, dtype=dtype, device=device)
1083110833
expected_unique_dim0 = torch.tensor([[[1., 1.],
1083210834
[0., 1.],
1083310835
[2., 1.],
@@ -10858,9 +10860,9 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
1085810860
device=device)
1085910861
expected_inverse_dim2 = torch.tensor([0, 1])
1086010862
expected_counts_dim2 = torch.tensor([1, 1])
10861-
expected_inverse_empty = torch.zeros(5, 0, dtype=dtype, device=device)
1086210863
expected_unique_empty = torch.tensor([], dtype=dtype, device=device)
10863-
expected_counts_empty = torch.tensor([], dtype=dtype, device=device)
10864+
expected_inverse_empty = torch.tensor([], dtype=torch.long, device=device)
10865+
expected_counts_empty = torch.tensor([], dtype=torch.long, device=device)
1086410866
# dim0
1086510867
x_unique = torch.unique(x, dim=0)
1086610868
self.assertEqual(expected_unique_dim0, x_unique)
@@ -10955,6 +10957,19 @@ def run_test(dtype=torch.float, device=torch.device('cpu')):
1095510957
self.assertEqual(expected_inverse_empty, x_inverse)
1095610958
self.assertEqual(expected_counts_empty, x_counts)
1095710959

10960+
# test not a well formed tensor
10961+
# Checking for runtime error, as this is the expected behaviour
10962+
self.assertRaises(RuntimeError, torch.unique,
10963+
x_ill_formed_empty,
10964+
return_inverse=True,
10965+
return_counts=True,
10966+
dim=1)
10967+
# test along dim2
10968+
self.assertRaises(RuntimeError, torch.unique,
10969+
x_ill_formed_empty_another,
10970+
return_inverse=True,
10971+
return_counts=True,
10972+
dim=2)
1095810973
# test consecutive version
1095910974
y = torch.tensor(
1096010975
[[0, 1],

0 commit comments

Comments
 (0)