DIM Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/AdaptiveAveragePooling.cpp lines 132–172

template <int64_t DIM>
std::vector<int64_t> get_output_shape(
    const Tensor& input,
    IntArrayRef output_size) {
  for (const auto i : c10::irange(1, input.dim())) {
    // Allow for empty batch.
    TORCH_CHECK(
        input.size(i) > 0,
        "adaptive_avg_pooling", DIM, "d(): ",
        "expected input to have non-empty spatial "
        "dimensions, but input has sizes ",
        input.sizes(),
        " with dimension ",
        i,
        " being empty");
  }

  TORCH_CHECK(
      (input.dim() == DIM + 1 || input.dim() == DIM + 2),
      "non-empty ",
      DIM + 1,
      "D or ",
      DIM + 2,
      "D (batch mode) tensor expected for input");

  /* Channels */
  const int64_t sizeC = input.size(-(DIM+1));

  std::vector<int64_t> output_shape;
  output_shape.reserve(input.dim());
  if (input.dim() == DIM + 2) {
    // Include Batch
    output_shape.push_back(input.size(0));
  }
  output_shape.push_back(sizeC);
  for (const auto size : output_size) {
    output_shape.push_back(size);
  }
  return output_shape;

}