Functions.h 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. #pragma once
  2. // ${generated_comment}
  3. #include <ATen/ATen.h>
  4. #include <ATen/core/functional.h>
  5. #include <ATen/TensorGeometry.h>
  6. #include "torch/csrc/autograd/function.h"
  7. #include "torch/csrc/autograd/variable.h"
  8. #include "torch/csrc/autograd/saved_variable.h"
  9. #include <torch/csrc/Export.h>
  10. #include <c10/core/SymIntArrayRef.h>
  11. namespace torch { namespace autograd { namespace generated {
  12. using at::Scalar;
  13. using at::Tensor;
  14. using at::IntArrayRef;
  15. using at::ArrayRef;
  16. using at::Type;
  17. using at::TensorGeometry;
  18. using at::ScalarType;
  19. using std::optional;
  20. using c10::fmap;
  21. inline std::vector<Tensor> unpack_list(at::ArrayRef<SavedVariable> xs, std::shared_ptr<Node> saved_for = nullptr) {
  22. // NB: we must explicitly do the conversion in the lambda, otherwise template
  23. // deduction will give a Tensor of Variable which is not convertible
  24. return fmap(xs, [&saved_for](const SavedVariable& x) {
  25. // TODO(crcrpar): Use `std::move(saved_for)` to avoid incrementing refcount, which would need refactoring.
  26. return static_cast<Tensor>(x.unpack(saved_for));
  27. });
  28. }
  29. inline c10::List<std::optional<Tensor>> unpack_opt_list(at::ArrayRef<SavedVariable> xs, std::shared_ptr<Node> saved_for = nullptr) {
  30. torch::List<std::optional<Tensor>> result;
  31. result.reserve(xs.size());
  32. for (const SavedVariable& v : xs) {
  33. auto var = v.unpack(saved_for);
  34. result.push_back(var.defined() ? std::optional<Tensor>(var) : ::std::nullopt);
  35. }
  36. return result;
  37. }
  38. using torch::autograd::TypeAndSize;
  39. ${autograd_function_declarations}
  40. }}} // namespace torch::autograd::generated