Skip to content

Commit f92edf7

Browse files
gchananfacebook-github-bot
gchanan
authored and
facebook-github-bot
committed
N-dimensional empty tensors: indexing, factories, reductions. (#9209)
Summary: This PR implements and tests N-dimensional empty tensors for indexing, factories, and reductions if compiled with -DUSE_TH_SIZE_ZERO_DIM. Still remaining to add: 1) TensorShape functions 2) Simple linear algebra functions (matrix multiply variants) 3) Other functions that operate over a dimension (but don't reduce). Pull Request resolved: #9209 Reviewed By: ezyang Differential Revision: D8751257 Pulled By: gchanan fbshipit-source-id: 2113374dc7af6caf31a99bf67b3893f130a29e23
1 parent 19ecb5f commit f92edf7

24 files changed

+567
-96
lines changed

‎aten/src/ATen/Declarations.cwrap‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1018,6 +1018,18 @@
10181018
return: real
10191019
arguments:
10201020
- THTensor* self
1021+
]]
1022+
[[
1023+
name: _th_all
1024+
types:
1025+
- Byte
1026+
variants:
1027+
- method
1028+
- function
1029+
backends:
1030+
- CPU
1031+
- CUDA
1032+
options:
10211033
- cname: logicalAnd
10221034
return: argument 0
10231035
scalar_check: self_->isScalar() || (keepdim == false && self_->dim() == 1)
@@ -1045,6 +1057,18 @@
10451057
return: real
10461058
arguments:
10471059
- THTensor* self
1060+
]]
1061+
[[
1062+
name: _th_any
1063+
types:
1064+
- Byte
1065+
variants:
1066+
- method
1067+
- function
1068+
backends:
1069+
- CPU
1070+
- CUDA
1071+
options:
10481072
- cname: logicalAny
10491073
return: argument 0
10501074
scalar_check: self_->isScalar() || (keepdim == false && self_->dim() == 1)
@@ -1641,6 +1665,18 @@
16411665
if_true: 0
16421666
if_false: 1
16431667
default: 0
1668+
]]
1669+
[[
1670+
name: _th_var
1671+
types:
1672+
- floating_point
1673+
backends:
1674+
- CPU
1675+
- CUDA
1676+
variants:
1677+
- method
1678+
- function
1679+
options:
16441680
- cname: var
16451681
return: argument 0
16461682
scalar_check: self_->isScalar() || (keepdim == false && self_->dim() == 1)
@@ -1676,6 +1712,18 @@
16761712
if_true: 0
16771713
if_false: 1
16781714
default: 0
1715+
]]
1716+
[[
1717+
name: _th_std
1718+
types:
1719+
- floating_point
1720+
backends:
1721+
- CPU
1722+
- CUDA
1723+
variants:
1724+
- method
1725+
- function
1726+
options:
16791727
- cname: std
16801728
return: argument 0
16811729
scalar_check: self_->isScalar() || (keepdim == false && self_->dim() == 1)
@@ -1711,7 +1759,7 @@
17111759
default: AS_REAL(2)
17121760
]]
17131761
[[
1714-
name: norm
1762+
name: _th_norm
17151763
types:
17161764
- floating_point
17171765
backends:

‎aten/src/ATen/ExpandUtils.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ std::vector<int64_t> infer_size(IntList a, IntList b) {
2121
") must match the size of tensor b (", sizeB,
2222
") at non-singleton dimension ", i);
2323

24-
expandedSizes[i] = std::max(sizeA, sizeB);
24+
// 1s map to the other size (even 0).
25+
expandedSizes[i] = sizeA == 1 ? sizeB : sizeA;
2526
}
2627

2728
return expandedSizes;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,6 @@
55
namespace at { namespace native {
66

77
Tensor pairwise_distance(const Tensor& x1, const Tensor& x2, double p, double eps, bool keepdim) {
8-
return norm(x1 - x2 + eps, p, 1, keepdim);
8+
return at::norm(x1 - x2 + eps, p, 1, keepdim);
99
}
1010
}} // namespace at::native

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ Tensor embedding_sparse_backward(
7171

7272
// check if all our grad come from padding_idx
7373
if (grad.numel() == 0) {
74+
// FIXME: USE_TH_SIZE_ZERO_DIM
7475
return sparse_type._sparse_coo_tensor_unsafe(indices_.type().tensor(),
7576
dense_type.tensor(), weight_size);
7677
}

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

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,13 @@ static std::vector<Tensor> expandByteTensors(const Tensor & self, TensorList ind
6969
}
7070
// Replace with nonzeros
7171
auto nonzero = index.nonzero();
72-
auto is_empty = nonzero.numel() == 0;
72+
#ifndef USE_TH_SIZE_ZERO_DIM
73+
auto special_empty = nonzero.numel() == 0;
74+
#else
75+
auto special_empty = false;
76+
#endif
7377
for (int64_t j = 0; j < index.dim(); j++) {
74-
if (is_empty) {
78+
if (special_empty) {
7579
// We can't call select on an empty tensor so we just create an empty
7680
// tensor.
7781
result.emplace_back(nonzero.type().tensor());
@@ -143,13 +147,15 @@ static Tensor unsqueezeN(const Tensor & src, int64_t before, int64_t after) {
143147
}
144148

145149
static Tensor wrapIndexOnce(const Tensor & index, int64_t dim, int64_t dim_size) {
146-
auto max_idx = index.max().toCLong();
147-
auto min_idx = index.min().toCLong();
148-
if (max_idx >= dim_size) {
149-
AT_ERROR("index ", max_idx, " is out of bounds for dimension ", dim, " with size ", dim_size);
150-
}
151-
if (min_idx < -dim_size) {
152-
AT_ERROR("index ", min_idx, " is out of bounds for dimension ", dim, " with size ", dim_size);
150+
if (index.numel() != 0) {
151+
auto max_idx = index.max().toCLong();
152+
auto min_idx = index.min().toCLong();
153+
if (max_idx >= dim_size) {
154+
AT_ERROR("index ", max_idx, " is out of bounds for dimension ", dim, " with size ", dim_size);
155+
}
156+
if (min_idx < -dim_size) {
157+
AT_ERROR("index ", min_idx, " is out of bounds for dimension ", dim, " with size ", dim_size);
158+
}
153159
}
154160
return index.remainder(dim_size);
155161
}
@@ -208,6 +214,7 @@ static Tensor computeLinearIndex(const Tensor & src, TensorList indices) {
208214
return linearIndex;
209215
}
210216

217+
#ifndef USE_TH_SIZE_ZERO_DIM
211218
static bool hasEmptyTensor(TensorList tensors) {
212219
for (auto& tensor : tensors) {
213220
if (tensor.defined() && tensor.numel() == 0) {
@@ -216,14 +223,17 @@ static bool hasEmptyTensor(TensorList tensors) {
216223
}
217224
return false;
218225
}
226+
#endif
219227

220228
static std::tuple<Tensor, Tensor> makeLinearIndex(Tensor self, TensorList orig) {
221229
checkIndexTensorTypes(orig);
222230
// first expand ByteTensor (boolean masks) into 1 or more LongTensors
223231
auto indices = expandByteTensors(self, orig);
232+
#ifndef USE_TH_SIZE_ZERO_DIM
224233
if (hasEmptyTensor(indices)) {
225234
return std::make_tuple(self, self.type().toScalarType(kLong).tensor());
226235
}
236+
#endif
227237
// next broadcast all index tensors together
228238
indices = expand_outplace(indices);
229239
// add missing null Tensors so that it matches self.dim()

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

Lines changed: 116 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44
#include "ATen/NativeFunctions.h"
55
#include "ATen/WrapDimUtils.h"
66
#include "ATen/WrapDimUtilsMulti.h"
7+
#include "ReduceOpsUtils.h"
78
#include "cpu/ReduceOpsKernel.h"
89

910
#include <algorithm>
1011
#include <functional>
12+
#include <limits>
1113
#include <numeric>
1214
#include <vector>
1315
#include <map>
@@ -94,10 +96,12 @@ static inline Tensor mean(const Tensor &self, optional<ScalarType> dtype) {
9496
"Can only calculate the mean of floating types. Got ",
9597
at::toString(scalarType),
9698
" instead.");
97-
Tensor result = at::native::sum(self);
98-
if (self.numel() > 0)
99-
result.div_(self.numel());
100-
return result;
99+
if (self.numel() > 0) {
100+
Tensor result = at::native::sum(self);
101+
return result.div_(self.numel());
102+
} else {
103+
return self.type().scalarTensor(std::numeric_limits<double>::quiet_NaN());
104+
}
101105
}
102106

103107
Tensor mean(const Tensor &self, ScalarType dtype) {
@@ -154,32 +158,6 @@ Tensor _prod_cpu(const Tensor &self) {
154158

155159
// DIM REDUCE #################################################################
156160

157-
static bool _dimreduce_return_trivial(Tensor &result, const Tensor &self,
158-
int64_t ident) {
159-
if (self.numel() == 1 && self.ndimension() == 0) {
160-
result.resize_({});
161-
result.fill_(self);
162-
return true;
163-
}
164-
// Return identity
165-
if (self.numel() == 0 && self.ndimension() == 1) {
166-
result.resize_({0});
167-
result.fill_(ident);
168-
return true;
169-
}
170-
return false;
171-
}
172-
173-
static Tensor &_dimreduce_setup(Tensor &result, const Tensor &self,
174-
int64_t dim) {
175-
IntList self_sizes = self.sizes();
176-
std::vector<int64_t> result_sizes;
177-
result_sizes.insert(result_sizes.end(), self_sizes.begin(), self_sizes.end());
178-
result_sizes[dim] = 1;
179-
result.resize_(result_sizes);
180-
return result;
181-
}
182-
183161
static inline Tensor &mean_out(Tensor &result, const Tensor &self, int64_t dim,
184162
bool keepdim, optional<ScalarType> dtype) {
185163
ScalarType scalarType = result.type().scalarType();
@@ -192,7 +170,12 @@ static inline Tensor &mean_out(Tensor &result, const Tensor &self, int64_t dim,
192170
result, self.toType(result.type().scalarType()), dim, keepdim);
193171
if (result.numel() > 0 && self.ndimension() > 0) {
194172
int64_t numel = self.size(dim);
195-
result.div_(numel);
173+
if (numel > 0) {
174+
result.div_(numel);
175+
} else {
176+
// NumPy equivalent
177+
result.fill_(std::numeric_limits<double>::quiet_NaN());
178+
}
196179
}
197180
return result;
198181
}
@@ -235,7 +218,7 @@ Tensor& sum_out(Tensor& result, const Tensor& self, IntList dim, ScalarType dtyp
235218
Tensor &_sum_out_cpu(Tensor &result, const Tensor &self, int64_t dim_,
236219
bool keepdim) {
237220
int64_t dim = maybe_wrap_dim(dim_, self.dim());
238-
if (_dimreduce_return_trivial(result, self, 0))
221+
if (_dimreduce_return_trivial(result, self, 0, dim, keepdim))
239222
return result;
240223
if (self.is_contiguous() && result.is_contiguous()) {
241224
_dimreduce_setup(result, self, dim);
@@ -273,7 +256,7 @@ Tensor& prod_out(Tensor& result, const Tensor& self, int64_t dim, ScalarType dty
273256
Tensor &_prod_out_cpu(Tensor &result, const Tensor &self, int64_t dim_,
274257
bool keepdim) {
275258
int64_t dim = maybe_wrap_dim(dim_, self.dim());
276-
if (_dimreduce_return_trivial(result, self, 1))
259+
if (_dimreduce_return_trivial(result, self, 1, dim, keepdim))
277260
return result;
278261
if (self.is_contiguous() && result.is_contiguous()) {
279262
_dimreduce_setup(result, self, dim);
@@ -294,7 +277,12 @@ static inline Tensor mean(const Tensor &self, int64_t dim, bool keepdim, optiona
294277
Tensor result = at::native::sum(self, dim, keepdim);
295278
if (result.numel() > 0 && self.ndimension() > 0) {
296279
int64_t numel = self.size(dim);
297-
result.div_(numel);
280+
if (numel > 0) {
281+
result.div_(numel);
282+
} else {
283+
// NumPy equivalent
284+
result.fill_(std::numeric_limits<double>::quiet_NaN());
285+
}
298286
}
299287
return result;
300288
}
@@ -357,10 +345,15 @@ Tensor _prod(const Tensor &self, int64_t dim_, bool keepdim) {
357345

358346
Tensor& logsumexp_out(Tensor& result, const Tensor &self, int64_t dim_, bool keepdim) {
359347
int64_t dim = maybe_wrap_dim(dim_, self.dim());
360-
auto maxes = at::max_values(self, dim, true);
361-
result = at::where((maxes == INFINITY).__or__(maxes == -INFINITY),
362-
maxes,
363-
maxes + at::log(at::sum(at::exp(self - maxes), dim, true)));
348+
// can't take max of empty tensor.
349+
if (self.numel() != 0) {
350+
auto maxes = at::max_values(self, dim, true);
351+
result = at::where((maxes == INFINITY).__or__(maxes == -INFINITY),
352+
maxes,
353+
maxes + at::log(at::sum(at::exp(self - maxes), dim, true)));
354+
} else {
355+
result = at::log(at::sum(at::exp(self), dim, true));
356+
}
364357
if (! keepdim)
365358
result.squeeze_(dim);
366359
return result;
@@ -588,4 +581,89 @@ Tensor& _sum_out(Tensor &result, const Tensor &self, IntList dims, bool keepdim)
588581
return reduce_multi_associative_out<_sum, _sum_out>(result, self, dims, keepdim);
589582
}
590583

584+
Tensor norm(const Tensor& self, Scalar p, int64_t dim, bool keepdim) {
585+
Tensor result = self.type().tensor();
586+
return at::native::norm_out(result, self, p, dim, keepdim);
587+
}
588+
589+
Tensor &norm_out(Tensor &result, const Tensor &self, Scalar p, int64_t dim, bool keepdim) {
590+
AT_CHECK(self.type().backend() == Backend::CPU || self.type().backend() == Backend::CUDA,
591+
"norm only supports CPU AND CUDA backend, got: ", at::toString(self.type().backend()));
592+
AT_CHECK(at::isFloatingType(self.type().scalarType()), "norm only supports floating-point dtypes");
593+
dim = maybe_wrap_dim(dim, self.dim());
594+
if (_dimreduce_return_trivial(result, self, 0, dim, keepdim)) {
595+
return result;
596+
} else {
597+
return at::_th_norm_out(result, self, p, dim, keepdim);
598+
}
599+
}
600+
601+
Tensor all(const Tensor& self, int64_t dim, bool keepdim) {
602+
Tensor result = self.type().tensor();
603+
return at::native::all_out(result, self, dim, keepdim);
604+
}
605+
606+
Tensor &all_out(Tensor &result, const Tensor &self, int64_t dim, bool keepdim) {
607+
AT_CHECK(self.type().backend() == Backend::CPU || self.type().backend() == Backend::CUDA,
608+
"all only supports CPU AND CUDA backend, got: ", at::toString(self.type().backend()));
609+
AT_CHECK(self.type().scalarType() == at::ScalarType::Byte, "all only supports torch.uint8 dtype");
610+
dim = maybe_wrap_dim(dim, self.dim());
611+
if (_dimreduce_return_trivial(result, self, 1, dim, keepdim)) {
612+
return result;
613+
} else {
614+
return at::_th_all_out(result, self, dim, keepdim);
615+
}
616+
}
617+
618+
Tensor any(const Tensor& self, int64_t dim, bool keepdim) {
619+
Tensor result = self.type().tensor();
620+
return at::native::any_out(result, self, dim, keepdim);
621+
}
622+
623+
Tensor &any_out(Tensor &result, const Tensor &self, int64_t dim, bool keepdim) {
624+
AT_CHECK(self.type().backend() == Backend::CPU || self.type().backend() == Backend::CUDA,
625+
"any only supports CPU AND CUDA backend, got: ", at::toString(self.type().backend()));
626+
AT_CHECK(self.type().scalarType() == at::ScalarType::Byte, "any only supports torch.uint8 dtype");
627+
dim = maybe_wrap_dim(dim, self.dim());
628+
if (_dimreduce_return_trivial(result, self, 0, dim, keepdim)) {
629+
return result;
630+
} else {
631+
return at::_th_any_out(result, self, dim, keepdim);
632+
}
633+
}
634+
635+
Tensor var(const Tensor& self, int64_t dim, bool unbiased, bool keepdim) {
636+
Tensor result = self.type().tensor();
637+
return at::native::var_out(result, self, dim, unbiased, keepdim);
638+
}
639+
640+
Tensor &var_out(Tensor &result, const Tensor &self, int64_t dim, bool unbiased, bool keepdim) {
641+
AT_CHECK(self.type().backend() == Backend::CPU || self.type().backend() == Backend::CUDA,
642+
"var only supports CPU AND CUDA backend, got: ", at::toString(self.type().backend()));
643+
AT_CHECK(at::isFloatingType(self.type().scalarType()), "var only supports floating-point dtypes");
644+
dim = maybe_wrap_dim(dim, self.dim());
645+
if (_dimreduce_return_trivial(result, self, std::numeric_limits<double>::quiet_NaN(), dim, keepdim)) {
646+
return result;
647+
} else {
648+
return at::_th_var_out(result, self, dim, unbiased, keepdim);
649+
}
650+
}
651+
652+
Tensor std(const Tensor& self, int64_t dim, bool unbiased, bool keepdim) {
653+
Tensor result = self.type().tensor();
654+
return at::native::std_out(result, self, dim, unbiased, keepdim);
655+
}
656+
657+
Tensor &std_out(Tensor &result, const Tensor &self, int64_t dim, bool unbiased, bool keepdim) {
658+
AT_CHECK(self.type().backend() == Backend::CPU || self.type().backend() == Backend::CUDA,
659+
"std only supports CPU AND CUDA backend, got: ", at::toString(self.type().backend()));
660+
AT_CHECK(at::isFloatingType(self.type().scalarType()), "std only supports floating-point dtypes");
661+
dim = maybe_wrap_dim(dim, self.dim());
662+
if (_dimreduce_return_trivial(result, self, std::numeric_limits<double>::quiet_NaN(), dim, keepdim)) {
663+
return result;
664+
} else {
665+
return at::_th_std_out(result, self, dim, unbiased, keepdim);
666+
}
667+
}
668+
591669
}} // namespace at::native

0 commit comments

Comments
 (0)