xref: /aosp_15_r20/external/pytorch/aten/src/ATen/WrapDimUtilsMulti.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)

#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