fractional_max_pool3d_backward_out_single_batch_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/FractionalMaxPool3d.cpp lines 259–287

template<typename scalar_t>
void fractional_max_pool3d_backward_out_single_batch_frame(
  scalar_t* gradInput,
  const scalar_t* gradOutput,
  const int64_t* indices,
  int64_t numPlanes,
  int64_t inputT, int64_t inputH, int64_t inputW,
  int64_t outputT, int64_t outputH, int64_t outputW) {

  at::parallel_for(0, numPlanes, 0, [&](int64_t start, int64_t end) {
    for (const auto plane : c10::irange(start, end)) {
      scalar_t* gradInputForPlane = gradInput + plane * inputT * inputH * inputW;
      const scalar_t* gradOutputForPlane = gradOutput +
                  plane * outputT * outputH * outputW;
      const int64_t* indicesForPlane = indices + plane * outputT * outputH * outputW;

      for (int64_t t = 0; t < outputT; ++t) {
        for (int64_t h = 0; h < outputH; ++h) {
          for (int64_t w = 0; w < outputW; ++w) {
            int64_t outputIndex = t * outputH * outputW + h * outputW + w;
            int64_t index = indicesForPlane[outputIndex];
            AT_ASSERT(index >= 0 && index < inputT * inputH * inputW);
            gradInputForPlane[index] += gradOutputForPlane[outputIndex];
          }
        }
      }
    }
  });
}