Compiler Explorer

Source code

#include <Eigen/Dense>
#ifdef __GNUC__
#define likely(x) __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)
#else
#define likely(x) (x)
#define unlikely(x) (x)
#endif

#ifndef STAN_MATH_MEMORY_STACK_ALLOC_HPP
#define STAN_MATH_MEMORY_STACK_ALLOC_HPP

// TODO(Bob): <cstddef> replaces this ifdef in C++11, until then this
//            is best we can do to get safe pointer casts to uints.
#include <stdint.h>
#include <cstdlib>
#include <cstddef>
#include <sstream>
#include <stdexcept>
#include <vector>

namespace stan {
namespace math {

/**
 * Return <code>true</code> if the specified pointer is aligned
 * on the number of bytes.
 *
 * This doesn't really make sense other than for powers of 2.
 *
 * @param ptr Pointer to test.
 * @param bytes_aligned Number of bytes of alignment required.
 * @return <code>true</code> if pointer is aligned.
 * @tparam Type of object to which pointer points.
 */
template <typename T>
bool is_aligned(T* ptr, unsigned int bytes_aligned) {
  return (reinterpret_cast<uintptr_t>(ptr) % bytes_aligned) == 0U;
}

namespace internal {
constexpr size_t DEFAULT_INITIAL_NBYTES = 1 << 16;  // 64KB

// FIXME: enforce alignment
// big fun to inline, but only called twice
inline char* eight_byte_aligned_malloc(size_t size) {
  char* ptr = static_cast<char*>(malloc(size));
  if (!ptr) {
    return ptr;  // malloc failed to alloc
  }
  if (!is_aligned(ptr, 8U)) {
    std::stringstream s;
    s << "invalid alignment to 8 bytes, ptr="
      << reinterpret_cast<uintptr_t>(ptr) << std::endl;
    throw std::runtime_error(s.str());
  }
  return ptr;
}
}  // namespace internal

/**
 * An instance of this class provides a memory pool through
 * which blocks of raw memory may be allocated and then collected
 * simultaneously.
 *
 * This class is useful in settings where large numbers of small
 * objects are allocated and then collected all at once.  This may
 * include objects whose destructors have no effect.
 *
 * Memory is allocated on a stack of blocks.  Each block allocated
 * is twice as large as the previous one.  The memory may be
 * recovered, with the blocks being reused, or all blocks may be
 * freed, resetting the stack of blocks to its original state.
 *
 * Alignment up to 8 byte boundaries guaranteed for the first malloc,
 * and after that it's up to the caller.  On 64-bit architectures,
 * all struct values should be padded to 8-byte boundaries if they
 * contain an 8-byte member or a virtual function.
 */
class stack_alloc {
 private:
  std::vector<char*> blocks_;  // storage for blocks,
                               // may be bigger than cur_block_
  std::vector<size_t> sizes_;  // could store initial & shift for others
  size_t cur_block_;           // index into blocks_ for next alloc
  char* cur_block_end_;        // ptr to cur_block_ptr_ + sizes_[cur_block_]
  char* next_loc_;             // ptr to next available spot in cur
                               // block
  // next three for keeping track of nested allocations on top of stack:
  std::vector<size_t> nested_cur_blocks_;
  std::vector<char*> nested_next_locs_;
  std::vector<char*> nested_cur_block_ends_;

/**
   * Moves us to the next block of memory, allocating that block
   * if necessary, and allocates len bytes of memory within that
   * block.
   *
   * @param len Number of bytes to allocate.
   * @return A pointer to the allocated memory.
   */
  inline char* move_to_next_block(size_t len) {
    char* result;
    ++cur_block_;
    // Find the next block (if any) containing at least len bytes.
    while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < len)) {
      ++cur_block_;
    }
    // Allocate a new block if necessary.
    if (unlikely(cur_block_ >= blocks_.size())) {
      // New block should be max(2*size of last block, len) bytes.
      size_t newsize = sizes_.back() * 2;
      if (newsize < len) {
        newsize = len;
      }
      blocks_.push_back(internal::eight_byte_aligned_malloc(newsize));
      if (!blocks_.back()) {
        throw std::bad_alloc();
      }
      sizes_.push_back(newsize);
    }
    result = blocks_[cur_block_];
    // Get the object's state back in order.
    next_loc_ = result + len;
    cur_block_end_ = result + sizes_[cur_block_];
    return result;
  }

public:
  /**
   * Construct a resizable stack allocator initially holding the
   * specified number of bytes.
   *
   * @param initial_nbytes Initial number of bytes for the
   * allocator.  Defaults to <code>(1 << 16) = 64KB</code> initial bytes.
   * @throws std::runtime_error if the underlying malloc is not 8-byte
   * aligned.
   */
  explicit stack_alloc(size_t initial_nbytes = internal::DEFAULT_INITIAL_NBYTES)
      : blocks_(1, internal::eight_byte_aligned_malloc(initial_nbytes)),
        sizes_(1, initial_nbytes),
        cur_block_(0),
        cur_block_end_(blocks_[0] + initial_nbytes),
        next_loc_(blocks_[0]) {
    if (!blocks_[0]) {
      throw std::bad_alloc();  // no msg allowed in bad_alloc ctor
    }
  }

/**
   * Destroy this memory allocator.
   *
   * This is implemented as a no-op as there is no destruction
   * required.
   */
  ~stack_alloc() {
    // free ALL blocks
    for (auto& block : blocks_) {
      if (block) {
        free(block);
      }
    }
  }

/**
   * Return a newly allocated block of memory of the appropriate
   * size managed by the stack allocator.
   *
   * The allocated pointer will be 8-byte aligned.
   *
   * This function may call C++'s <code>malloc()</code> function,
   * with any exceptions percolated through this function.
   *
   * @param len Number of bytes to allocate.
   * @return A pointer to the allocated memory.
   */
  inline void* alloc(size_t len) {
    // Typically, just return and increment the next location.
    char* result = next_loc_;
    next_loc_ += len;
    // Occasionally, we have to switch blocks.
    if (unlikely(next_loc_ >= cur_block_end_)) {
      result = move_to_next_block(len);
    }
    return reinterpret_cast<void*>(result);
  }

/**
   * Allocate an array on the arena of the specified size to hold
   * values of the specified template parameter type.
   *
   * @tparam T type of entries in allocated array.
   * @param[in] n size of array to allocate.
   * @return new array allocated on the arena.
   */
  template <typename T>
  inline T* alloc_array(size_t n) {
    return static_cast<T*>(alloc(n * sizeof(T)));
  }

/**
   * Recover all the memory used by the stack allocator.  The stack
   * of memory blocks allocated so far will be available for further
   * allocations.  To free memory back to the system, use the
   * function free_all().
   */
  inline void recover_all() {
    cur_block_ = 0;
    next_loc_ = blocks_[0];
    cur_block_end_ = next_loc_ + sizes_[0];
  }

/**
   * Store current positions before doing nested operation so can
   * recover back to start.
   */
  inline void start_nested() {
    nested_cur_blocks_.push_back(cur_block_);
    nested_next_locs_.push_back(next_loc_);
    nested_cur_block_ends_.push_back(cur_block_end_);
  }

/**
   * recover memory back to the last start_nested call.
   */
  inline void recover_nested() {
    if (unlikely(nested_cur_blocks_.empty())) {
      recover_all();
    }

cur_block_ = nested_cur_blocks_.back();
    nested_cur_blocks_.pop_back();

next_loc_ = nested_next_locs_.back();
    nested_next_locs_.pop_back();

cur_block_end_ = nested_cur_block_ends_.back();
    nested_cur_block_ends_.pop_back();
  }

/**
   * Free all memory used by the stack allocator other than the
   * initial block allocation back to the system.  Note:  the
   * destructor will free all memory.
   */
  inline void free_all() {
    // frees all BUT the first (index 0) block
    for (size_t i = 1; i < blocks_.size(); ++i) {
      if (blocks_[i]) {
        free(blocks_[i]);
      }
    }
    sizes_.resize(1);
    blocks_.resize(1);
    recover_all();
  }

/**
   * Return number of bytes allocated to this instance by the heap.
   * This is not the same as the number of bytes allocated through
   * calls to memalloc_.  The latter number is not calculatable
   * because space is wasted at the end of blocks if the next
   * alloc request doesn't fit.  (Perhaps we could trim down to
   * what is actually used?)
   *
   * @return number of bytes allocated to this instance
   */
  inline size_t bytes_allocated() const {
    size_t sum = 0;
    for (size_t i = 0; i <= cur_block_; ++i) {
      sum += sizes_[i];
    }
    return sum;
  }

/**
   * Indicates whether the memory in the pointer
   * is in the stack.
   *
   * @param[in] ptr memory location
   * @return true if the pointer is in the stack,
   *    false otherwise.
   */
  inline bool in_stack(const void* ptr) const {
    for (size_t i = 0; i < cur_block_; ++i) {
      if (ptr >= blocks_[i] && ptr < blocks_[i] + sizes_[i]) {
        return true;
      }
    }
    if (ptr >= blocks_[cur_block_] && ptr < next_loc_) {
      return true;
    }
    return false;
  }
};

}  // namespace math
}  // namespace stan
#endif

#include <vector>
#ifndef STAN_MATH_REV_CORE_AUTODIFFSTACKSTORAGE_HPP
#define STAN_MATH_REV_CORE_AUTODIFFSTACKSTORAGE_HPP
#include <vector>

namespace stan {
namespace math {

// Internal macro used to modify global pointer definition to the
// global AD instance.
#ifdef STAN_THREADS
// Whenever STAN_THREADS is set a TLS keyword is used. For reasons
// explained below we use the GNU compiler extension __thread if
// supported by the compiler while the generic thread_local C++11
// keyword is used otherwise.
#ifdef __GNUC__
#define STAN_THREADS_DEF __thread
#else
#define STAN_THREADS_DEF thread_local
#endif
#else
// In case STAN_THREADS is not set, then no modifier is needed.
#define STAN_THREADS_DEF
#endif

/**
 * This struct always provides access to the autodiff stack using
 * the singleton pattern. Read warnings below!
 *
 * The singleton <code>instance_</code> is a global static pointer,
 * which is thread local (TLS) if the STAN_THREADS preprocess variable
 * is defined.
 *
 * The use of a pointer is motivated by performance reasons for the
 * threading case. When a TLS is used, initialization with a constant
 * expression at compile time is required for fast access to the
 * TLS. As the autodiff storage struct is non-POD, its initialization
 * is a dynamic expression at compile time. These dynamic expressions
 * are wrapped, in the TLS case, by a TLS wrapper function which slows
 * down its access. Using a pointer instead allows to initialize at
 * compile time to <code>nullptr</code>, which is a compile time
 * constant. In this case, the compiler avoids the use of a TLS
 * wrapper function.
 *
 * For performance reasons we use the __thread keyword on compilers
 * which support it. The __thread keyword is a GNU compiler-specific
 * (gcc, clang, Intel) extension which requires initialization with a
 * compile time constant expression. The C++11 keyword thread_local
 * does allow for constant and dynamic initialization of the
 * TLS. Thus, only the __thread keyword guarantees that constant
 * initialization and its implied speedup, is used.
 *
 * The initialization of the AD instance at run-time is handled by the
 * lifetime of a AutodiffStackSingleton object. More specifically, the
 * first instance of the AutodiffStackSingleton object will initialize
 * the AD instance and take ownership (it is the only one instance
 * with the private member own_instance_ being true). Thus, whenever
 * the first instance of the AutodiffStackSingleton object gets
 * destructed, the AD tape will be destructed as well. Within
 * stan-math the initialization of the AD instance for the main thread
 * of the program is handled by instantiating the singleton once in
 * the init_chainablestack.hpp file. Whenever STAN_THREADS is defined
 * then all created child threads must instantiate a
 * AutodiffStackSingleton object within the child thread before
 * accessing the AD system in order to initialize the TLS AD tape
 * within the child thread.
 *
 * The design of a globally held (optionally TLS) pointer, which is
 * globally initialized, allows the compiler to apply necessary
 * inlining to get maximal performance. However, the design suffers
 * from "the static init order fiasco"[0]. Whenever the static init
 * order fiasco occurs, the C++ client of the library may instantiate
 * an AutodiffStackSingleton object at the adequate code position prior
 * to any AD tape access to ensure proper initialization order. In
 * exchange, we get a more performant singleton pattern with automatic
 * initialization of the AD stack for the main thread. There has been
 * some discussion on earlier designs using the Meyer singleton
 * approach; see [1] and [2] and the discussions those PRs link to as
 * well.
 *
 * [0] https://isocpp.org/wiki/faq/ctors#static-init-order
 * [1] https://github.com/stan-dev/math/pull/840
 * [2] https://github.com/stan-dev/math/pull/826
 * [3]
 * http://discourse.mc-stan.org/t/potentially-dropping-support-for-older-versions-of-apples-version-of-clang/3780/
 */
template <typename ChainableT, typename ChainableAllocT>
struct AutodiffStackSingleton {
  using AutodiffStackSingleton_t
      = AutodiffStackSingleton<ChainableT, ChainableAllocT>;

AutodiffStackSingleton() : own_instance_(init()) {}
  ~AutodiffStackSingleton() {
    if (own_instance_) {
      delete instance_;
      instance_ = nullptr;
    }
  }

struct AutodiffStackStorage {
    AutodiffStackStorage &operator=(const AutodiffStackStorage &) = delete;

std::vector<ChainableT> var_stack_;
    std::vector<ChainableT> var_nochain_stack_;
    std::vector<ChainableAllocT *> var_alloc_stack_;
    stack_alloc memalloc_;

// nested positions
    std::vector<size_t> nested_var_stack_sizes_;
    std::vector<size_t> nested_var_nochain_stack_sizes_;
    std::vector<size_t> nested_var_alloc_stack_starts_;
  };

explicit AutodiffStackSingleton(AutodiffStackSingleton_t const &) = delete;
  AutodiffStackSingleton &operator=(const AutodiffStackSingleton_t &) = delete;

static STAN_THREADS_DEF AutodiffStackStorage *instance_;

private:
  static bool init() {
    static STAN_THREADS_DEF bool is_initialized = false;
    if (!is_initialized) {
      is_initialized = true;
      instance_ = new AutodiffStackStorage();
      return true;
    }
    if (!instance_) {
      is_initialized = true;
      instance_ = new AutodiffStackStorage();
      return true;
    }
    return false;
  }

bool own_instance_;
};

template <typename ChainableT, typename ChainableAllocT>
STAN_THREADS_DEF
    typename AutodiffStackSingleton<ChainableT,
                                    ChainableAllocT>::AutodiffStackStorage
        *AutodiffStackSingleton<ChainableT, ChainableAllocT>::instance_;

}  // namespace math
}  // namespace stan
#endif

#ifndef STAN_MATH_REV_CORE_CHAINABLE_ALLOC_HPP
#define STAN_MATH_REV_CORE_CHAINABLE_ALLOC_HPP

#ifndef STAN_MATH_REV_CORE_CHAINABLESTACK_HPP
#define STAN_MATH_REV_CORE_CHAINABLESTACK_HPP

#include <boost/variant2/variant.hpp>
namespace stan {
namespace math {

template <typename T, typename = void>
class vari_value;
class vari_base;
class chainable_alloc;

using vari_variant
    = boost::variant2::variant<vari_value<float>*,
       vari_value<Eigen::Matrix<double, -1, -1>>*, vari_value<long double>*,
       vari_value<Eigen::Matrix<double, -1, 1>>*, vari_value<Eigen::Matrix<double, 1, -1>>*, vari_value<double>*>;

using ChainableStack = AutodiffStackSingleton<vari_variant, chainable_alloc>;

}  // namespace math
}  // namespace stan
#endif

namespace stan {
namespace math {

/**
 * A chainable_alloc is an object which is constructed and
 * destructed normally but the memory lifespan is managed along
 * with the arena allocator for the gradient calculation.  A
 * chainable_alloc instance must be created with a call to
 * operator new for memory management.
 */
class chainable_alloc {
 public:
  chainable_alloc() {
    ChainableStack::instance_->var_alloc_stack_.push_back(this);
  }
};

}  // namespace math
}  // namespace stan
#endif
#ifndef STAN_MATH_REV_CORE_INIT_CHAINABLESTACK_HPP
#define STAN_MATH_REV_CORE_INIT_CHAINABLESTACK_HPP

#include <tbb/task_scheduler_observer.h>

#include <mutex>
#include <unordered_map>
#include <utility>
#include <thread>
#include <tuple>

namespace stan {
namespace math {

/**
 * TBB observer object which is a callback hook called whenever the
 * TBB scheduler adds a new thread to the TBB managed threadpool. This
 * hook ensures that each worker thread has an initialized AD tape
 * ready for use.
 *
 * Refer to https://software.intel.com/en-us/node/506314 for details
 * on the observer concept.
 */
class ad_tape_observer final : public tbb::task_scheduler_observer {
  using stack_ptr = std::unique_ptr<ChainableStack>;
  using ad_map = std::unordered_map<std::thread::id, stack_ptr>;

public:
  ad_tape_observer() : tbb::task_scheduler_observer(), thread_tape_map_() {
    on_scheduler_entry(true);  // register current process
    observe(true);             // activates the observer
  }

~ad_tape_observer() { observe(false); }

inline void on_scheduler_entry(bool worker) {
    std::lock_guard<std::mutex> thread_tape_map_lock(thread_tape_map_mutex_);
    const std::thread::id thread_id = std::this_thread::get_id();
    if (thread_tape_map_.find(thread_id) == thread_tape_map_.end()) {
      ad_map::iterator insert_elem;
      bool status = false;
      std::tie(insert_elem, status)
          = thread_tape_map_.emplace(ad_map::value_type{thread_id, nullptr});
      insert_elem->second = stack_ptr(new ChainableStack());
    }
  }

inline void on_scheduler_exit(bool worker) {
    std::lock_guard<std::mutex> thread_tape_map_lock(thread_tape_map_mutex_);
    auto elem = thread_tape_map_.find(std::this_thread::get_id());
    if (elem != thread_tape_map_.end()) {
      thread_tape_map_.erase(elem);
    }
  }

private:
  ad_map thread_tape_map_;
  std::mutex thread_tape_map_mutex_;
};

namespace {

ad_tape_observer global_observer;

}  // namespace
}  // namespace math
}  // namespace stan

#endif

#ifndef STAN_MATH_REV_CORE_VARI_HPP
#define STAN_MATH_REV_CORE_VARI_HPP

#include <ostream>
#include <type_traits>

namespace stan {
namespace math {

// forward declaration of var
template <typename T, typename>
class var_value;

/**
 * The variable implementation base class.
 *
 * This class is complete (not abstract) and may be used for
 * constants.
 *
 * A variable implementation is constructed with a constant
 * value. It also stores the adjoint for storing the partial
 * derivative with respect to the root of the derivative tree.
 *
 * The chain() method applies the chain rule. Concrete extensions
 * of this class will represent base variables or the result
 * of operations such as addition or subtraction. These extended
 * classes will store operand variables and propagate derivative
 * information via an implementation of chain().
 */
template <typename T, typename>
class vari_value;

template <typename T>
class vari_value<T, std::enable_if_t<std::is_floating_point<T>::value>> {
 private:
  template <typename, typename>
  friend class var_value;

public:
  using Scalar = T;
  using value_type = Scalar;
  /**
   * The value of this variable.
   */
  const Scalar val_;

/**
   * The adjoint of this variable, which is the partial derivative
   * of this variable with respect to the root variable.
   */
  Scalar adj_;

/**
   * Construct a variable implementation from a value.  The
   * adjoint is initialized to zero.
   *
   * All constructed variables are added to the stack.  Variables
   * should be constructed before variables on which they depend
   * to insure proper partial derivative propagation.  During
   * derivative propagation, the chain() method of each variable
   * will be called in the reverse order of construction.
   *
   * @tparam S an Arithmetic type.
   * @param x Value of the constructed variable.
   */
  template <typename S,
            std::enable_if_t<std::is_convertible<S&, Scalar>::value>* = nullptr>
  vari_value(S x) noexcept : val_(x), adj_(0.0) {  // NOLINT
    ChainableStack::instance_->var_stack_.emplace_back(this);
  }

/**
   * Construct a variable implementation from a value.  The
   *  adjoint is initialized to zero and if `stacked` is `false` this vari
   *  will be moved to the nochain stack s.t. it's chain method will not be
   *  called when calling `grad()`.
   *
   * All constructed variables are added to the stack.  Variables
   *  should be constructed before variables on which they depend
   *  to insure proper partial derivative propagation.  During
   *  derivative propagation, the chain() method of each variable
   *  will be called in the reverse order of construction.
   *
   * @tparam S an Arithmetic type.
   * @param x Value of the constructed variable.
   * @param stacked If false will put this this vari on the nochain stack so
   * that its `chain()` method is not called.
   */
  template <typename S,
            std::enable_if_t<std::is_convertible<S&, Scalar>::value>* = nullptr>
  vari_value(S x, bool stacked) noexcept : val_(x), adj_(0.0) {
    if (stacked) {
      ChainableStack::instance_->var_stack_.emplace_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.emplace_back(this);
    }
  }

/**
   * Constructor from vari_value
   * @tparam S An arithmetic type
   * @param x A vari_value
   */
  template <typename S,
            std::enable_if_t<std::is_arithmetic<S>::value>* = nullptr>
  vari_value(const vari_value<S>& x) noexcept : val_(x.val_), adj_(x.adj_) {
    ChainableStack::instance_->var_stack_.emplace_back(this);
  }
  /**
   * Constructor from vari_value
   * @tparam S An arithmetic type
   * @param x A vari_value
   */
  template <typename S,
            std::enable_if_t<std::is_arithmetic<S>::value>* = nullptr>
  vari_value(vari_value<S>&& x) noexcept : val_(x.val_), adj_(x.adj_) {
    ChainableStack::instance_->var_stack_.emplace_back(this);
  }

/**
   * Initialize the adjoint for this (dependent) variable to 1.
   * This operation is applied to the dependent variable before
   * propagating derivatives, setting the derivative of the
   * result with respect to itself to be 1.
   */
  inline void init_dependent() noexcept { adj_ = 1.0; }

/**
   * Set the adjoint value of this variable to 0.  This is used to
   * reset adjoints before propagating derivatives again (for
   * example in a Jacobian calculation).
   */
  inline void set_zero_adjoint() noexcept { adj_ = 0.0; }

virtual void chain() {}

/**
   * Insertion operator for vari. Prints the current value and
   * the adjoint value.
   *
   * @param os [in, out] ostream to modify
   * @param v [in] vari object to print.
   *
   * @return The modified ostream.
   */
  friend std::ostream& operator<<(std::ostream& os, const vari_value<T>* v) {
    return os << v->val_ << ":" << v->adj_;
  }

/**
   * Allocate memory from the underlying memory pool.  This memory is
   * is managed as a whole externally.
   *
   * Warning: Classes should not be allocated with this operator
   * if they have non-trivial destructors.
   *
   * @param nbytes Number of bytes to allocate.
   * @return Pointer to allocated bytes.
   */
  static inline void* operator new(size_t nbytes) noexcept {
    return ChainableStack::instance_->memalloc_.alloc(nbytes);
  }

/**
   * Delete a pointer from the underlying memory pool.
   *
   * This no-op implementation enables a subclass to throw
   * exceptions in its constructor.  An exception thrown in the
   * constructor of a subclass will result in an error being
   * raised, which is in turn caught and calls delete().
   *
   * See the discussion of "plugging the memory leak" in:
   *   http://www.parashift.com/c++-faq/memory-pools.html
   */
  static inline void operator delete(void* /* ignore arg */) noexcept { /* no op */
  }
};

// For backwards compatability the default is double
using vari = vari_value<double>;

struct vari_printer {
  std::ostream& o_;
  int i_{0};
  vari_printer(std::ostream& o, int i) : o_(o), i_(i) {}
  template <typename T>
  inline void operator()(T*& x) const {
    o_ << i_ << "  " << x << "  " << x->val_ << " : " << x->adj_ << std::endl;
  }
};

template <typename T>
struct is_eigen : std::false_type {};

template <typename T>
struct is_eigen<Eigen::Matrix<T, -1, -1>> : std::true_type {};

template <typename T>
struct is_eigen<Eigen::Matrix<T, 1, -1>> : std::true_type {};

template <typename T>
struct is_eigen<Eigen::Matrix<T, -1, 1>> : std::true_type {};

template <typename T>
struct value_type {
    using type = T;
};

template <typename T>
struct value_type<Eigen::Matrix<T, -1, -1>> {
    using type = T;
};

template <typename T>
struct value_type<Eigen::Matrix<T, 1, -1>> {
    using type = T;
};

template <typename T>
struct value_type<Eigen::Matrix<T, -1, 1>> {
    using type = T;
};

template <typename T>
using value_type_t = typename value_type<std::decay_t<T>>::type;

template <typename T>
class vari_value<T, std::enable_if_t<is_eigen<T>::value>> {
 private:
  template <typename, typename>
  friend class var_value;
  using eigen_scalar = value_type_t<T>;
  value_type_t<T>* val_mem_;
  value_type_t<T>* adj_mem_;

public:
  using Scalar = T;
  using value_type = T;
  using PlainObject = typename T::PlainObject;
  using eigen_map = Eigen::Map<PlainObject>;
  const Eigen::Index rows_;
  const Eigen::Index cols_;
  const Eigen::Index size_;
  const Eigen::Index RowsAtCompileTime = T::RowsAtCompileTime;
  const Eigen::Index ColsAtCompileTime = T::ColsAtCompileTime;
  /**
   * The value of this variable.
   */
  eigen_map val_;

/**
   * The adjoint of this variable, which is the partial derivative
   * of this variable with respect to the root variable.
   */
  eigen_map adj_;

/**
   * Construct a variable implementation from a value.  The
   * adjoint is initialized to zero.
   *
   * All constructed variables are added to the stack.  Variables
   * should be constructed before variables on which they depend
   * to insure proper partial derivative propagation.  During
   * derivative propagation, the chain() method of each variable
   * will be called in the reverse order of construction.
   *
   * @param x Value of the constructed variable.
   */
  explicit vari_value(const T& x)
      : val_mem_(ChainableStack::instance_->memalloc_.alloc_array<eigen_scalar>(
            x.size())),
        adj_mem_(ChainableStack::instance_->memalloc_.alloc_array<eigen_scalar>(
            x.size())),
        rows_(x.rows()),
        cols_(x.cols()),
        size_(x.size()),
        val_(val_mem_, x.rows(), x.cols()),
        adj_(adj_mem_, x.rows(), x.cols()) {
    val_ = x;
    adj_.setZero();
    ChainableStack::instance_->var_stack_.push_back(this);
  }

vari_value(const T& x, bool stacked)
      : val_mem_(ChainableStack::instance_->memalloc_.alloc_array<eigen_scalar>(
            x.size())),
        adj_mem_(ChainableStack::instance_->memalloc_.alloc_array<eigen_scalar>(
            x.size())),
        rows_(x.rows()),
        cols_(x.cols()),
        size_(x.size()),
        val_(val_mem_, x.rows(), x.cols()),
        adj_(adj_mem_, x.rows(), x.cols()) {
    val_ = x;
    adj_.setZero();
    if (stacked) {
      ChainableStack::instance_->var_stack_.push_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.push_back(this);
    }
  }

const Eigen::Index rows() const { return rows_; }
  const Eigen::Index cols() const { return cols_; }
  const Eigen::Index size() const { return size_; }

/**
   * Set the adjoint value of this variable to 0.  This is used to
   * reset adjoints before propagating derivatives again (for
   * example in a Jacobian calculation).
   */
  void set_zero_adjoint() { adj_.setZero(); }

/**
   * Allocate memory from the underlying memory pool.  This memory is
   * is managed as a whole externally.
   *
   * Warning: Classes should not be allocated with this operator
   * if they have non-trivial destructors.
   *
   * @param nbytes Number of bytes to allocate.
   * @return Pointer to allocated bytes.
   */
  static inline void* operator new(size_t nbytes) {
    return ChainableStack::instance_->memalloc_.alloc(nbytes);
  }

#ifndef STAN_MATH_REV_CORE_SET_ZERO_ALL_ADJOINTS_HPP
#define STAN_MATH_REV_CORE_SET_ZERO_ALL_ADJOINTS_HPP

namespace stan {
namespace math {

struct zero_adj {
  template <typename T, std::enable_if_t<std::is_arithmetic<T>::value>* = nullptr>
  EIGEN_STRONG_INLINE void operator()(vari_value<T>*& x) const {
      x->adj_ = 0.0;
  }
  template <typename T, std::enable_if_t<is_eigen<T>::value>* = nullptr>
  EIGEN_STRONG_INLINE void operator()(vari_value<T>*& x) const {
      x->adj_.fill(0.0);
  }

};

template <typename T>
using require_t = std::enable_if_t<T::value>;

template <bool B>
using bool_constant = std::integral_constant<bool, B>;

template <size_t N, typename = void>
struct visit_nt_impl;

template <size_t N>
struct visit_nt_impl<N, require_t<bool_constant<(N == 0)>>> {
  template <typename Variant, typename Visitor>
  static constexpr auto apply(Variant &&var, Visitor &&vis) noexcept {
        if (N == var.index()) {
        // If this check isnt there the compiler will generate
        // exception code, this stops that
        return std::forward<Visitor>(vis)(
            boost::variant2::get<N>(std::forward<Variant>(var)));
        } 
  }
};
template <size_t N>
struct visit_nt_impl<N, require_t<bool_constant<(N > 0)>>> {
  template <typename Variant, typename Visitor>
  static constexpr auto apply(Variant &&var, Visitor &&vis) noexcept {
        if (var.index() == N) {
        return std::forward<Visitor>(vis)(
            boost::variant2::get<N>(std::forward<Variant>(var)));
        } else {
        return visit_nt_impl<N - 1>::apply(std::forward<Variant>(var),
                                std::forward<Visitor>(vis));
        }
    __builtin_unreachable();
  }
};

template <typename Variant, typename Visitor>
constexpr auto visit_nt(Variant &&var, Visitor &&vis) noexcept {
  constexpr auto sizer = boost::variant2::variant_size<Variant>::value; 
  return visit_nt_impl<sizer - 1>::apply(std::forward<Variant>(var), std::forward<Visitor>(vis));
}

/**
 * Reset all adjoint values in the stack to zero.
 */
static EIGEN_STRONG_INLINE void set_zero_all_adjoints() noexcept {
  for (auto& x : ChainableStack::instance_->var_stack_) {
    visit_nt(x, zero_adj());
  }
  for (auto& x : ChainableStack::instance_->var_nochain_stack_) {
    visit_nt(x, zero_adj());
  }
}

}  // namespace math
}  // namespace stan
#endif

int main() {
    using stan::math::vari;
    std::vector<vari*> A;
    for (int i = 0; i < 10; i++) {
      A.push_back(new vari(static_cast<double>(i)));
    }
    stan::math::set_zero_all_adjoints();
    return 1;
}