File size: 1,075 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
#pragma once
#include <ATen/WrapDimUtils.h>
#include <c10/core/TensorImpl.h>
#include <c10/util/irange.h>
#include <bitset>
#include <sstream>
namespace at {
// This is in an extra file to work around strange interaction of
// bitset on Windows with operator overloading
constexpr size_t dim_bitset_size = 64;
inline std::bitset<dim_bitset_size> dim_list_to_bitset(
OptionalIntArrayRef opt_dims,
size_t ndims) {
TORCH_CHECK(
ndims <= dim_bitset_size,
"only tensors with up to ",
dim_bitset_size,
" dims are supported");
std::bitset<dim_bitset_size> seen;
if (opt_dims.has_value()) {
auto dims = opt_dims.value();
for (const auto i : c10::irange(dims.size())) {
size_t dim = maybe_wrap_dim(dims[i], static_cast<int64_t>(ndims));
TORCH_CHECK(
!seen[dim],
"dim ",
dim,
" appears multiple times in the list of dims");
seen[dim] = true;
}
} else {
for (size_t dim = 0; dim < ndims; dim++) {
seen[dim] = true;
}
}
return seen;
}
} // namespace at
|