cpu_pixel_unshuffle_channels_last Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/PixelShuffleKernel.cpp lines 154–194

template <typename scalar_t>
void cpu_pixel_unshuffle_channels_last(
    TensorBase& output,
    const TensorBase& input,
    int64_t downscale_factor) {
  TORCH_CHECK(input.ndimension() == 4,
              "pixel unshuffle with channels last format supports tensors with 4 dims");
  auto input_data = input.const_data_ptr<scalar_t>();
  auto output_data = output.data_ptr<scalar_t>();

  int64_t nbatch = input.size(0);
  int64_t sub_channels = input.size(1);
  int64_t height = input.size(2) / downscale_factor;
  int64_t width = input.size(3) / downscale_factor;
  int64_t channels = sub_channels * downscale_factor * downscale_factor;
  int64_t numel = input.numel();
  int64_t S = downscale_factor;

  // input strides
  int64_t stride_n = height * width * channels;
  int64_t stride_h = S * width * S * sub_channels;
  int64_t stride_s1 = width * S * sub_channels;
  int64_t stride_w = S * sub_channels;
  int64_t stride_s2 = sub_channels;
  int64_t stride_c = 1;

  // input tensor shape of [n, h, s1, w, s2, c]
  // output tensor shape of [n, h, w, c, s1, s2]
  at::parallel_for(0, numel, 0, [&](int64_t begin, int64_t end) {
    int64_t n{0}, h{0}, w{0}, c{0}, s1{0}, s2{0};
    data_index_init(begin, n, nbatch, h, height, w, width, c, sub_channels, s1, S, s2, S);

    for (const auto i : c10::irange(begin, end)) {
      int64_t input_offset = n * stride_n + h * stride_h + s1 * stride_s1 +
          w * stride_w + s2 * stride_s2 + c * stride_c;
      output_data[i] = c10::load(&input_data[input_offset]);

      data_index_step(n, nbatch, h, height, w, width, c, sub_channels, s1, S, s2, S);
    }
  });
}