Function Transpose

oneflow/core/framework/tensor_methods.cpp:735–776

}

Maybe<Tensor> Transpose(const std::shared_ptr<Tensor>& input, const std::vector<int32_t>& permute) {
  const auto& shape = input->shape();
  const auto& strides = JUST(input->stride());
  const int64_t ndim = shape->NumAxes();
  int64_t storage_offset = JUST(JUST(input->AsLocalTensor())->storage_offset());

  CHECK_EQ_OR_RETURN(permute.size(), ndim)
      << "permute size should be equal to input tensor's ndim, but got " << permute.size();
  auto positive_perm = permute;
  for (auto i = 0; i < positive_perm.size(); i++) {
    positive_perm[i] = JUST(maybe_wrap_dim(positive_perm[i], ndim));
  }

  DimVector target_dims(ndim);
  Stride stride(ndim);
  for (int i = 0; i < ndim; ++i) {
    target_dims[i] = shape->At(permute[i]);
    stride[i] = strides->at(permute[i]);
  }

  auto output = JUST(BasicView(input, Shape(target_dims), 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> {
      std::vector<int32_t> grad_perm;
      grad_perm.resize(ndim);
      for (int i = 0; i < ndim; ++i) { grad_perm[permute[i]] = i; }
      autograd::AutoGradMode mode(create_graph);
      CHECK_EQ_OR_RETURN(out_grads.size(), 1)
          << "out grad size should be 1, but got " << out_grads.size();
      in_grads->resize(1);
      (*in_grads)[0] = JUST(functional::Transpose(out_grads[0], grad_perm));
      return Maybe<void>::Ok();
    };
    backward_fn->status = []() { return true; };
    TensorTuple outputs{output};
    JUST(GetThreadLocalAutogradEngine()->AddNode("view::transpose_backward", backward_fn, {input},
                                                 &outputs));
  }
  return output;
}

Maybe<Tensor> UnfoldTensor(const std::shared_ptr<Tensor>& input, const int32_t dimension,
                           const int32_t size, const int32_t step) {