WrapDimUtilsMulti.h 802 B

123456789101112131415161718192021222324252627
  1. #pragma once
  2. #include <c10/core/TensorImpl.h>
  3. #include <c10/util/irange.h>
  4. #include <ATen/WrapDimUtils.h>
  5. #include <sstream>
  6. #include <bitset>
  7. namespace at {
  8. // This is in an extra file to work around strange interaction of
  9. // bitset on Windows with operator overloading
  10. constexpr size_t dim_bitset_size = 64;
  11. static inline std::bitset<dim_bitset_size> dim_list_to_bitset(IntArrayRef dims, int64_t ndims) {
  12. TORCH_CHECK(ndims <= (int64_t) dim_bitset_size, "only tensors with up to ", dim_bitset_size, " dims are supported");
  13. std::bitset<dim_bitset_size> seen;
  14. for (const auto i : c10::irange(dims.size())) {
  15. size_t dim = maybe_wrap_dim(dims[i], ndims);
  16. TORCH_CHECK(!seen[dim], "dim ", dim, " appears multiple times in the list of dims");
  17. seen[dim] = true;
  18. }
  19. return seen;
  20. }
  21. }