Function Narrow

oneflow/core/framework/tensor_methods.cpp:515–556

}

Maybe<Tensor> Narrow(const std::shared_ptr<Tensor>& input, const int64_t dim, const int64_t start,
                     const int64_t length) {
  const auto& shape = input->shape();
  const auto& strides = JUST(input->stride());
  const int64_t ndim = shape->NumAxes();
  DimVector dim_vec;
  dim_vec.insert(dim_vec.end(), shape->dim_vec().cbegin(), shape->dim_vec().cbegin() + dim);
  dim_vec.insert(dim_vec.end(), length);
  dim_vec.insert(dim_vec.end(), shape->dim_vec().cbegin() + dim + 1, shape->dim_vec().end());

  int64_t storage_offset = JUST(JUST(input->AsLocalTensor())->storage_offset());
  Shape target_shape(dim_vec);

  Stride stride(ndim);
  for (int i = 0; i < ndim; ++i) {
    stride[i] = strides->at(i);
    if (dim == i) { storage_offset += start * strides->at(i); }
  }

  auto output = JUST(BasicView(input, target_shape, stride, storage_offset));
  if (autograd::GradMode::is_enabled() && input->requires_grad()) {
    auto backward_fn = std::make_shared<BackwardFunction>();
    backward_fn->body = [=](const TensorTuple& out_grads, TensorTuple* in_grads,
                            bool create_graph) -> Maybe<void> {
      autograd::AutoGradMode mode(create_graph);
      CHECK_EQ_OR_RETURN(out_grads.size(), 1)
          << "out grad size should be 1, but got " << out_grads.size();
      auto like =
          JUST(functional::Empty(Shape(input->shape()->dim_vec()), input->dtype(),
                                 JUST(input->device()), /*requires_grad=*/input->requires_grad(),
                                 /*pin_memory=*/false));
      in_grads->resize(1);
      (*in_grads)[0] = JUST(functional::NarrowGrad(out_grads[0], like, dim, start, length));
      return Maybe<void>::Ok();
    };
    backward_fn->status = []() { return true; };
    TensorTuple outputs{output};
    JUST(GetThreadLocalAutogradEngine()->AddNode("view::narrow_backward", backward_fn, {input},
                                                 &outputs));
  }
  return output;
}

Maybe<Tensor> AsStridedGrad(const std::shared_ptr<one::Tensor>& dy,
                            const std::shared_ptr<one::Tensor>& input,