Function AsStridedGrad

oneflow/core/framework/tensor_methods.cpp:558–679

}

Maybe<Tensor> AsStridedGrad(const std::shared_ptr<one::Tensor>& dy,
                            const std::shared_ptr<one::Tensor>& input,
                            const std::vector<int64_t>& sizes, const std::vector<int64_t>& strides,
                            const int64_t storage_offset) {
  CHECK_OR_RETURN(input->is_local()) << "input must be local tensor.";
  // reference: torch/csrc/autograd/FunctionsManual.cpp
  const size_t odim = dy->ndim();
  std::vector<int64_t> out_sizes_, out_strides_;
  out_sizes_.reserve(odim);
  out_strides_.reserve(odim);
  auto grad = dy;
  for (int64_t i = odim - 1; i >= 0; i--) {
    auto size_i = sizes[i];
    auto stride_i = strides[i];
    if (size_i == 0) {
      return functional::Constant(*dy->shape(), 0, grad->dtype(), JUST(grad->device()));
    } else if (size_i == 1) {
      grad = JUST(functional::Squeeze(grad, std::vector<int32_t>{int(i)}));
    } else if (stride_i == 0) {
      grad = JUST(functional::ReduceSum(grad, std::vector<int32_t>{int(i)}, false, NullOpt));
    } else {
      out_sizes_.insert(out_sizes_.begin(), size_i);
      out_strides_.insert(out_strides_.begin(), stride_i);
    }
  }

  // Step (2)~(4) for the algorithm in NOTE [ Detecting Memory Overlap Within A
  // Strided Tensor ]
  //              on output geometry
  const bool out_maybe_overlap = IsOverlappingMemorys(out_sizes_, out_strides_);

  // For input geometry,
  //   check for size 0 dimensions,
  //   skip size 1 dimensions,
  // Step (0)~(1) for the algorithm in NOTE [ Detecting Memory Overlap Within A
  // Strided Tensor ]
  //              on input geometry
  auto idim = input->ndim();
  std::vector<int64_t> inp_sizes(input->shape()->begin(), input->shape()->end());
  std::vector<int64_t> inp_strides(JUST(input->stride())->begin(), JUST(input->stride())->end());
  std::vector<int64_t> inp_sizes_, inp_strides_;
  inp_sizes_.reserve(idim);
  inp_strides_.reserve(idim);
  for (int64_t i = idim - 1; i >= 0; i--) {
    auto size_i = inp_sizes[i];
    auto stride_i = inp_strides[i];
    if (size_i == 0) {
      return functional::Constant(*input->shape(), 0, grad->dtype(), JUST(grad->device()));
    } else if (size_i != 1) {
      inp_sizes_.insert(inp_sizes_.begin(), size_i);
      inp_strides_.insert(inp_strides_.begin(), stride_i);
    }
  }
  // Step (1)~(4) for the algorithm in NOTE [ Detecting Memory Overlap Within A
  // Strided Tensor ]
  //              on input geometry
  const bool inp_maybe_overlap = IsOverlappingMemorys(inp_sizes_, inp_strides_);