code: #include <torch/csrc/autograd/forward_grad.h>
  
  namespace torch { namespace autograd {
  
  namespace {
      // See discussion in forward_grad.h for why these are global variables and not
      // thread local
  
      // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
      static std::mutex all_forward_levels_mutex_;
      // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
      static uint64_t next_forward_idx_ = 0;
      // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
      static std::vector<std::shared_ptr<ForwardADLevel>> all_forward_levels_;
  
      const static at::Tensor singleton_undefined_tensor;
  
      // Temporary flag to disable forward mode
      // TODO(alband) remove these when perf issues are solved
      // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
      static bool is_forward_grad_enabled = false;
  }
  
  uint64_t ForwardADLevel::get_next_idx() {
      std::lock_guard<std::mutex> lock(all_forward_levels_mutex_);
      TORCH_CHECK(next_forward_idx_ == 0, "Nested forward mode AD is not supported at the moment");
      auto new_index = next_forward_idx_++;
      TORCH_INTERNAL_ASSERT(new_index == all_forward_levels_.size());
      all_forward_levels_.push_back(std::make_shared<ForwardADLevel>(new_index));
      return new_index;
  }
  
  void ForwardADLevel::release_idx(uint64_t idx) {
      std::lock_guard<std::mutex> lock(all_forward_levels_mutex_);
      TORCH_CHECK(idx == all_forward_levels_.size() - 1, "Exiting a forward AD level that is not the "
                  "last that was created is not support. Ensure they are released in the reverse "
                  "order they were created.");
      TORCH_CHECK(idx >= 0, "No forward AD level was created so you cannot exit it.");
      next_forward_idx_--;
      all_forward_levels_.pop_back();
  
  }
  std::shared_ptr<ForwardADLevel> ForwardADLevel::get_by_idx(uint64_t idx) {
      std::lock_guard<std::mutex> lock(all_forward_levels_mutex_);
      TORCH_CHECK(idx < all_forward_levels_.size(), "Trying to access a forward AD level with an invalid index. "
                  "This index was either not created or is already deleted.");
      return all_forward_levels_[idx];
  }
  
  std::shared_ptr<ForwardADLevel> ForwardADLevel::try_get_by_idx(uint64_t idx) {
      std::lock_guard<std::mutex> lock(all_forward_levels_mutex_);
      if (idx < all_forward_levels_.size()) {
          return all_forward_levels_[idx];
      } else {
          return nullptr;
      }
  }
  
  ForwardADLevel::~ForwardADLevel() {
      std::lock_guard<std::mutex> lock(mutex_);
      auto it = grads_.begin();
      while (it != grads_.end()) {
          // Warning this will lock *it mutex
          // This is ok as this function is the *only* one to call back into another class's method.
          (*it)->reset(idx_, /* update_level */ false);
          it = grads_.erase(it);
      }
  }
  
  const at::Tensor& ForwardGrad::value(uint64_t level) const {
      std::lock_guard<std::mutex> lock(mutex_);
      const auto& it = content_.find(level);
      return it == content_.end() ? singleton_undefined_tensor : (*it).second;
  }
  
  const at::Tensor& ForwardGrad::undef_grad() {
      return singleton_undefined_tensor;
  }
  
  // Temporary functions to disable forward AD
  // TODO(alband) remove these when perf issues are solved
  bool isForwardADEnabled() {
      return is_forward_grad_enabled;
  }
  
  void setForwardADEnabled(bool value) {
      is_forward_grad_enabled = value;
  }
  
  }} // namespace torch::autograd
  