#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once

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

namespace at {
// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
class TORCH_API OptionalTensorRef {
 public:
  OptionalTensorRef() = default;

  ~OptionalTensorRef() {
    ref_.unsafeReleaseTensorImpl();
  }

  OptionalTensorRef(const TensorBase& src C10_LIFETIMEBOUND)
      : ref_(Tensor::unsafe_borrow_t{}, src) {
    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src.defined());
  }

  OptionalTensorRef(const OptionalTensorRef& rhs)
      : ref_(Tensor::unsafe_borrow_t{}, rhs.ref_) {}

  OptionalTensorRef(OptionalTensorRef&& rhs) = default;
  OptionalTensorRef& operator=(OptionalTensorRef rhs) {
    std::swap(ref_, rhs.ref_);
    return *this;
  }

  bool has_value() const {
    return ref_.defined();
  }

  const Tensor& getTensorRef() const & C10_LIFETIMEBOUND {
    return ref_;
  }

  const Tensor& operator*() const & C10_LIFETIMEBOUND {
    return ref_;
  }

  const Tensor* operator->() const & C10_LIFETIMEBOUND {
    return &ref_;
  }

  operator bool() const {
    return ref_.defined();
  }

 private:
  Tensor ref_;
};

// Use to convert a TensorBase (that may be undefined) to an at::Tensor
// without bumping refcount.
class TORCH_API TensorRef {
 public:
  ~TensorRef() {
    ref_.unsafeReleaseTensorImpl();
  }

  TensorRef(const TensorBase& src C10_LIFETIMEBOUND)
      : ref_(Tensor::unsafe_borrow_t{}, src) {}
  TensorRef(TensorRef&& other) = default;
  TensorRef(const TensorRef&) = default;
  TensorRef& operator=(const TensorRef&) = default;
  TensorRef& operator=(TensorRef&&) = default;

  const Tensor& operator*() const & C10_LIFETIMEBOUND {
    return ref_;
  }
 private:
  Tensor ref_;
};

template <typename T>
unsigned Tensor::register_hook(T&& hook) const {
  using R = std::invoke_result_t<T, Tensor>;
  if constexpr (std::is_same_v<R, Tensor>) {
    return _register_hook([fn = std::forward<T>(hook)](const TensorBase& grad_base) {
      TensorRef grad(grad_base);
      Tensor ret = fn(*grad);
      return TensorBase(std::move(ret));
    });
  } else {
    static_assert(std::is_void_v<R>, "hook must return Tensor or void");
    return _register_hook([fn = std::forward<T>(hook)](const TensorBase& grad_base) {
      TensorRef grad(grad_base);
      fn(*grad);
      return Tensor();
    });
  }
}

} // namespace at

#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
