variable.h

#pragma once

#include <torch/csrc/utils/python_stub.h>

#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/csrc/autograd/edge.h>
#include <torch/csrc/autograd/function_hook.h>
#include <torch/csrc/autograd/variable_version.h>

#include <ATen/ATen.h>
#include <c10/util/Exception.h>

#include <list>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace torch { namespace autograd {

struct Function;


struct TORCH_API Variable : public at::Tensor {
  Variable() = default;

  // Factory Functions
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  // NOTE: These factory functions have to be friends to access the
  // `Variable::Impl`. As a side effect, it allows us to keep them in the class.

  friend Variable make_variable_view(
      Variable base,
      at::Tensor data,
      bool is_differentiable,
      bool allow_tensor_metadata_change,
      Edge gradient_edge);

  friend Variable make_variable(
      at::Tensor data,
      bool requires_grad,
      bool allow_tensor_metadata_change);

  friend Variable make_variable_consuming(
      at::Tensor data,
      bool requires_grad,
      bool allow_tensor_metadata_change);

  friend Variable make_variable(
      at::Tensor data,
      Edge gradient_edge,
      bool allow_tensor_metadata_change);

  // Tensor Conversions
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  // "Downcasts" a `Tensor` into a `Variable`. Only call this on tensors you
  // know are Variables.
  /*implicit*/ Variable(at::Tensor const& rhs) : at::Tensor(rhs) {
    AT_CHECK(
        is_variable() || !defined(),
        "Tensor that was converted to Variable was not actually a Variable");
  }

  /*implicit*/ Variable(at::Tensor&& rhs)
      : at::Tensor(std::move(rhs)) {
    AT_CHECK(
        is_variable() || !defined(),
        "Tensor that was converted to Variable was not actually a Variable");
  }

  // NOTE: Assignment operators to Tensor come for free from the constructors.

  const at::Tensor& data() const noexcept;
  at::Tensor& data() noexcept;

  // Gradient Function and Edges
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  const std::shared_ptr<Function>& grad_fn() const;

  Function* grad_fn_unsafe() const;

  void set_grad_accumulator(std::weak_ptr<Function> grad_accumulator);

  std::shared_ptr<Function> try_get_grad_accumulator() const;

  std::shared_ptr<Function> grad_accumulator() const;

  Edge gradient_edge() const {
    // If grad_fn is null (as is the case for a leaf node), we instead
    // interpret the gradient function to be a gradient accumulator, which will
    // accumulate its inputs into the grad property of the variable. These
    // nodes get suppressed in some situations, see "suppress gradient
    // accumulation" below. Note that only variables which have `requires_grad =
    // True` can have gradient accumulators.
    if (const auto& gradient = grad_fn()) {
      return Edge(gradient, output_nr());
    } else {
      return Edge(grad_accumulator(), 0);
    }
  }

  Variable detach() const;

  void detach_();

  void backward(
      c10::optional<Tensor> gradient,
      bool keep_graph,
      bool create_graph) const;

  void set_data(const at::Tensor &new_data);

  void set_gradient_edge(Edge edge) noexcept;

  uint32_t output_nr() const noexcept;

  bool is_leaf() const noexcept;

  // Versions
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  void bump_version() noexcept;
  void set_version_counter(const VariableVersion& version_counter) noexcept;

  const VariableVersion& version_counter() const noexcept;

  uint32_t current_version() const noexcept;

  // Autograd Graph Interaction
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  void rebase_history(Edge gradient_edge);

  // Hooks
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  void add_hook(std::shared_ptr<FunctionPreHook> hook);
  const std::vector<std::shared_ptr<FunctionPreHook>>& hooks() const noexcept;
  void clear_hooks();

  // View Variables
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  bool is_view() const noexcept;

  const Variable& base() const;

  // Miscellaneous
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  void set_name(const std::string& name);
  const std::string& name() const noexcept;

  PyObject* pyobj() const noexcept;
  void set_pyobj(PyObject* pyobj) noexcept;

  struct AutogradMeta;
  Variable::AutogradMeta* get_autograd_meta() const noexcept;

 private:
  struct Impl;
  struct DifferentiableViewImpl;
  struct DifferentiableViewMeta;

  // Private Methods
  //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  Variable(c10::intrusive_ptr<Variable::Impl> self);
  Impl* get() const;
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
//                            Variable::AutogradMeta
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


struct TORCH_API Variable::AutogradMeta : public c10::AutogradMetaInterface {
  std::string name;

  Variable grad_;
  std::shared_ptr<Function> grad_fn_;
  std::weak_ptr<Function> grad_accumulator_;

  VariableVersion version_counter_;
  std::vector<std::shared_ptr<FunctionPreHook>> hooks_;

  // Only meaningful on leaf variables (must be false otherwise)
  bool requires_grad_;

  bool is_view_;

  // The "output number" of this variable; e.g., if this variable
  // was the second output of a function, then output_nr == 1.
  // We use this to make sure we can setup the backwards trace
  // correctly when this variable is passed to another function.
  uint32_t output_nr_;
  PyObject* pyobj_ = nullptr; // weak reference

  // Mutex to ensure that concurrent read operations that modify internal
  // state are still thread-safe. Used by grad_fn() and
  // grad_accumulator().
  std::mutex mutex_;

  void set_requires_grad(bool requires_grad, at::TensorImpl* self_impl) override {
    AT_CHECK(
      !requires_grad || at::isFloatingType(at::typeMetaToScalarType(self_impl->dtype())),
      "Only Tensors of floating point dtype can require gradients");
    requires_grad_ = requires_grad;
  }

  bool requires_grad() const override {
    return requires_grad_ || grad_fn_;
  }

  Variable& grad() override {
    return grad_;
  }

  const Variable& grad() const override {
    return grad_;
  }
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
//                        Variable::DifferentiableViewMeta
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

struct TORCH_API Variable::DifferentiableViewMeta : public Variable::AutogradMeta {
  Variable base_;

  uint32_t attr_version;

  bool requires_grad() const override {
    return requires_grad_ || grad_fn_ || (is_view_ && base_.requires_grad());
  }
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
//                            Variable::Impl
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

struct TORCH_API Variable::Impl : public at::TensorImpl {
  explicit Impl(
      at::Tensor data,
      std::unique_ptr<Variable::AutogradMeta> autograd_meta,
      bool requires_grad = false,
      Edge gradient_edge = Edge());

  ~Impl() override;

  int64_t numel() const override;
  at::IntArrayRef sizes() const override;
  at::IntArrayRef strides() const override;
  bool is_contiguous() const override;
  int64_t size(int64_t d) const override;
  int64_t stride(int64_t d) const override;
  void resize_dim(int64_t ndim) override;
  void set_size(int64_t dim, int64_t new_size) override;
  void set_stride(int64_t dim, int64_t new_stride) override;
  void set_storage_offset(int64_t storage_offset) override;

  int64_t dim() const override;
  bool has_storage() const override;
  const at::Storage& storage() const override;
  void* slow_data() const override;

  void set_data(const at::Tensor &new_data);

  void release_resources() override;

  Variable::AutogradMeta* get_autograd_meta() const {
    return static_cast<Variable::AutogradMeta*>(autograd_meta());
  }

  int64_t storage_offset() const override;

  at::Tensor data_;

 private:
  int64_t get_device_slow() const override;
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
//                     Variable::DifferentiableViewImpl
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

struct TORCH_API Variable::DifferentiableViewImpl : public Variable::Impl {
  DifferentiableViewImpl(
    Variable base,
    at::Tensor data,
    Edge gradient_edge,
    std::unique_ptr<Variable::DifferentiableViewMeta> autograd_meta);

  void release_resources() override;
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
//                        Variable Implementation
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

// Factory Functions
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


// See NOTE [ Autograd View Variables ] for details.
inline Variable make_variable_view(
    Variable base,
    at::Tensor data,
    bool is_differentiable = true,
    bool allow_tensor_metadata_change = true,
    Edge gradient_edge = Edge()) {
  if (data.defined()) {
    if (is_differentiable) {
      auto data_impl_copy = data.getIntrusivePtr()->shallow_copy_and_detach();
      data_impl_copy->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
      auto data_copy = at::Tensor(data_impl_copy);
      auto diff_view_meta = c10::guts::make_unique<Variable::DifferentiableViewMeta>();
      return Variable(c10::make_intrusive<Variable::DifferentiableViewImpl>(
              std::move(base), std::move(data_copy), std::move(gradient_edge), std::move(diff_view_meta)));
    } else {
      auto data_impl_copy = data.getIntrusivePtr()->shallow_copy_and_detach();
      data_impl_copy->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
      auto data_copy = at::Tensor(data_impl_copy);
      auto autograd_meta = c10::guts::make_unique<Variable::AutogradMeta>();
      auto var = Variable(c10::make_intrusive<Variable::Impl>(
              std::move(data_copy), std::move(autograd_meta), false, std::move(gradient_edge)));
      var.set_version_counter(base.version_counter());
      return var;
    }
  }
  return Variable();
}

inline Variable make_variable(
    at::Tensor data,
    bool requires_grad = false,
    bool allow_tensor_metadata_change = true) {
  AT_CHECK(
      !data.is_variable(),
      "Must not create a new variable from a variable, use its .data()");
  if (data.defined()) {
    auto data_impl_copy = data.getIntrusivePtr()->shallow_copy_and_detach();
    data_impl_copy->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
    auto data_copy = at::Tensor(data_impl_copy);
    auto autograd_meta = c10::guts::make_unique<Variable::AutogradMeta>();
    return Variable(c10::make_intrusive<Variable::Impl>(data_copy, std::move(autograd_meta), requires_grad));
  }
  return Variable();
}

inline Variable make_variable_consuming(
    at::Tensor data,
    bool requires_grad = false,
    bool allow_tensor_metadata_change = true) {
  AT_CHECK(
      !data.is_variable(),
      "Must not create a new variable from a variable, use its .data()");
  if (data.defined()) {
    AT_ASSERT(data.getIntrusivePtr().use_count() == 1);
    data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
    auto autograd_meta = c10::guts::make_unique<Variable::AutogradMeta>();
    return Variable(c10::make_intrusive<Variable::Impl>(std::move(data), std::move(autograd_meta), requires_grad));
  }
  return Variable();
}

inline Variable make_variable(
    at::Tensor data,
    Edge gradient_edge,
    bool allow_tensor_metadata_change = true) {
  AT_CHECK(
      !data.is_variable(),
      "Must not create a new variable from a variable, use its .data()");
  if (data.defined()) {
    auto data_impl_copy = data.getIntrusivePtr()->shallow_copy_and_detach();
    data_impl_copy->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
    auto data_copy = at::Tensor(data_impl_copy);
    auto autograd_meta = c10::guts::make_unique<Variable::AutogradMeta>();
    return Variable(c10::make_intrusive<Variable::Impl>(data_copy, std::move(autograd_meta), false, std::move(gradient_edge)));
  }
  return Variable();
}

// Tensor Conversion
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline Variable& as_variable_ref(at::Tensor& tensor) {
  AT_CHECK(
      tensor.is_variable(),
      "Attempted to cast a Tensor to a Variable, but "
      "the dynamic type of the value is not Variable.");
  return static_cast<Variable&>(tensor);
}

inline const Variable& as_variable_ref(const at::Tensor& tensor) {
  AT_CHECK(
      tensor.is_variable(),
      "Attempted to cast a Tensor to a Variable, but "
      "the dynamic type of the value is not Variable.");
  return static_cast<const Variable&>(tensor);
}

inline const at::Tensor& Variable::data() const noexcept {
  return get()->data_;
}

inline at::Tensor& Variable::data() noexcept {
  return get()->data_;
}

// Gradient Function and Edges
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline Function* Variable::grad_fn_unsafe() const {
  return get_autograd_meta()->grad_fn_.get();
}

inline void Variable::set_grad_accumulator(
    std::weak_ptr<Function> grad_accumulator) {
  get_autograd_meta()->grad_accumulator_ = std::move(grad_accumulator);
}

inline std::shared_ptr<Function> Variable::try_get_grad_accumulator() const {
  return get_autograd_meta()->grad_accumulator_.lock();
}

inline Variable Variable::detach() const {
  auto var = make_variable_view(*this, get()->data_, /*is_differentiable=*/false, /*allow_tensor_metadata_change=*/false, Edge());
  return var;
}

inline void Variable::set_data(const at::Tensor &new_data) {
  get()->set_data(new_data);
}

inline void Variable::set_gradient_edge(Edge edge) noexcept {
  get_autograd_meta()->grad_fn_ = std::move(edge.function);
  get_autograd_meta()->output_nr_ = edge.input_nr;
}

inline uint32_t Variable::output_nr() const noexcept {
  return get_autograd_meta()->output_nr_;
}

inline bool Variable::is_leaf() const noexcept {
  return get_autograd_meta()->grad_fn_ == nullptr;
}

// Versions
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline void Variable::set_version_counter(
    const VariableVersion& version_counter) noexcept {
  get_autograd_meta()->version_counter_ = version_counter;
}

inline void Variable::bump_version() noexcept {
  get_autograd_meta()->version_counter_.bump();
}

inline uint32_t Variable::current_version() const noexcept {
  return get_autograd_meta()->version_counter_.current_version();
}

inline const VariableVersion& Variable::version_counter() const noexcept {
  return get_autograd_meta()->version_counter_;
}

// Hooks
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline void Variable::add_hook(std::shared_ptr<FunctionPreHook> hook) {
  get_autograd_meta()->hooks_.push_back(std::move(hook));
}

inline const std::vector<std::shared_ptr<FunctionPreHook>>& Variable::hooks()
    const noexcept {
  return get_autograd_meta()->hooks_;
}

inline void Variable::clear_hooks() {
  get_autograd_meta()->hooks_.clear();
}

// View Variables
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline bool Variable::is_view() const noexcept {
  return get_autograd_meta()->is_view_;
}

inline const Variable& Variable::base() const {
  if (is_view()) {
    auto diff_view_meta = static_cast<Variable::DifferentiableViewMeta*>(get_autograd_meta());
    return diff_view_meta->base_;
  } else {
    throw std::runtime_error("Can't get base of non-view Variable");
  }
}

// Miscellaneous
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline void Variable::set_name(const std::string& name) {
  get_autograd_meta()->name = name;
}

inline const std::string& Variable::name() const noexcept {
  return get_autograd_meta()->name;
}

inline void Variable::set_pyobj(PyObject* pyobj) noexcept {
  get_autograd_meta()->pyobj_ = pyobj;
}

inline PyObject* Variable::pyobj() const noexcept {
  return get_autograd_meta()->pyobj_;
}

inline Variable::AutogradMeta* Variable::get_autograd_meta() const noexcept {
  return get()->get_autograd_meta();
}

// Private Methods
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

inline Variable::Variable(c10::intrusive_ptr<Variable::Impl> self)
    : at::Tensor(std::move(self)) {}

inline Variable::Impl* Variable::get() const {
  AT_CHECK(defined(), "Called Variable::get() on an undefined Variable");
  return static_cast<Variable::Impl*>(impl_.get());
}
}} // namespace torch::autograd