variable_matrix.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <algorithm>
#include <array>
#include <concepts>
#include <cstddef>
#include <initializer_list>
#include <iterator>
#include <span>
#include <utility>
#include <vector>
 
#include <Eigen/Core>
#include <Eigen/LU>
#include <Eigen/QR>
#include <gch/small_vector.hpp>
 
#include "sleipnir/autodiff/expression_graph.hpp"
#include "sleipnir/autodiff/sleipnir_base.hpp"
#include "sleipnir/autodiff/slice.hpp"
#include "sleipnir/autodiff/variable.hpp"
#include "sleipnir/autodiff/variable_block.hpp"
#include "sleipnir/util/assert.hpp"
#include "sleipnir/util/concepts.hpp"
#include "sleipnir/util/empty.hpp"
#include "sleipnir/util/function_ref.hpp"
#include "sleipnir/util/symbol_exports.hpp"
 
namespace slp {
 
template <typename Scalar_>
class VariableMatrix : public SleipnirBase {
 public:
  using Scalar = Scalar_;
 
  VariableMatrix() = default;
 
  explicit VariableMatrix(int rows) : VariableMatrix{rows, 1} {}
 
  VariableMatrix(int rows, int cols) : m_rows{rows}, m_cols{cols} {
    m_storage.reserve(rows * cols);
    for (int index = 0; index < rows * cols; ++index) {
      m_storage.emplace_back();
    }
  }
 
  VariableMatrix(detail::empty_t, int rows, int cols)
      : m_rows{rows}, m_cols{cols} {
    m_storage.reserve(rows * cols);
    for (int index = 0; index < rows * cols; ++index) {
      m_storage.emplace_back(nullptr);
    }
  }
 
  VariableMatrix(
      std::initializer_list<std::initializer_list<Variable<Scalar>>> list) {
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert all column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      slp_assert(static_cast<int>(row.size()) == m_cols);
    }
 
    m_storage.reserve(rows() * cols());
    for (const auto& row : list) {
      std::ranges::copy(row, std::back_inserter(m_storage));
    }
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const std::vector<std::vector<Scalar>>& list) {
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert all column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      slp_assert(static_cast<int>(row.size()) == m_cols);
    }
 
    m_storage.reserve(rows() * cols());
    for (const auto& row : list) {
      std::ranges::copy(row, std::back_inserter(m_storage));
    }
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const std::vector<std::vector<Variable<Scalar>>>& list) {
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert all column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      slp_assert(static_cast<int>(row.size()) == m_cols);
    }
 
    m_storage.reserve(rows() * cols());
    for (const auto& row : list) {
      std::ranges::copy(row, std::back_inserter(m_storage));
    }
  }
 
  template <typename Derived>
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const Eigen::MatrixBase<Derived>& values)
      : m_rows{static_cast<int>(values.rows())},
        m_cols{static_cast<int>(values.cols())} {
    m_storage.reserve(values.rows() * values.cols());
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        m_storage.emplace_back(values[row, col]);
      }
    }
  }
 
  template <typename Derived>
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const Eigen::DiagonalBase<Derived>& values)
      : m_rows{static_cast<int>(values.rows())},
        m_cols{static_cast<int>(values.cols())} {
    m_storage.reserve(values.rows() * values.cols());
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        if (row == col) {
          m_storage.emplace_back(values.diagonal()[row]);
        } else {
          m_storage.emplace_back(Scalar(0));
        }
      }
    }
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const Variable<Scalar>& variable) : m_rows{1}, m_cols{1} {
    m_storage.emplace_back(variable);
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(Variable<Scalar>&& variable) : m_rows{1}, m_cols{1} {
    m_storage.emplace_back(std::move(variable));
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const VariableBlock<VariableMatrix>& values)
      : m_rows{values.rows()}, m_cols{values.cols()} {
    m_storage.reserve(rows() * cols());
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        m_storage.emplace_back(values[row, col]);
      }
    }
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  VariableMatrix(const VariableBlock<const VariableMatrix>& values)
      : m_rows{values.rows()}, m_cols{values.cols()} {
    m_storage.reserve(rows() * cols());
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        m_storage.emplace_back(values[row, col]);
      }
    }
  }
 
  explicit VariableMatrix(std::span<const Variable<Scalar>> values)
      : m_rows{static_cast<int>(values.size())}, m_cols{1} {
    m_storage.reserve(rows() * cols());
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        m_storage.emplace_back(values[row * cols() + col]);
      }
    }
  }
 
  VariableMatrix(std::span<const Variable<Scalar>> values, int rows, int cols)
      : m_rows{rows}, m_cols{cols} {
    slp_assert(static_cast<int>(values.size()) == rows * cols);
    m_storage.reserve(rows * cols);
    for (int row = 0; row < rows; ++row) {
      for (int col = 0; col < cols; ++col) {
        m_storage.emplace_back(values[row * cols + col]);
      }
    }
  }
 
  template <typename Derived>
  VariableMatrix& operator=(const Eigen::MatrixBase<Derived>& values) {
    slp_assert(rows() == values.rows() && cols() == values.cols());
 
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        (*this)[row, col] = values[row, col];
      }
    }
 
    return *this;
  }
 
  VariableMatrix& operator=(ScalarLike auto value) {
    slp_assert(rows() == 1 && cols() == 1);
 
    (*this)[0, 0] = value;
 
    return *this;
  }
 
  template <typename Derived>
    requires std::same_as<typename Derived::Scalar, Scalar>
  void set_value(const Eigen::MatrixBase<Derived>& values) {
    slp_assert(rows() == values.rows() && cols() == values.cols());
 
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        (*this)[row, col].set_value(values[row, col]);
      }
    }
  }
 
  Variable<Scalar>& operator[](int row, int col) {
    slp_assert(row >= 0 && row < rows());
    slp_assert(col >= 0 && col < cols());
    return m_storage[row * cols() + col];
  }
 
  const Variable<Scalar>& operator[](int row, int col) const {
    slp_assert(row >= 0 && row < rows());
    slp_assert(col >= 0 && col < cols());
    return m_storage[row * cols() + col];
  }
 
  Variable<Scalar>& operator[](int index) {
    slp_assert(index >= 0 && index < rows() * cols());
    return m_storage[index];
  }
 
  const Variable<Scalar>& operator[](int index) const {
    slp_assert(index >= 0 && index < rows() * cols());
    return m_storage[index];
  }
 
  VariableBlock<VariableMatrix> block(int row_offset, int col_offset,
                                      int block_rows, int block_cols) {
    slp_assert(row_offset >= 0 && row_offset <= rows());
    slp_assert(col_offset >= 0 && col_offset <= cols());
    slp_assert(block_rows >= 0 && block_rows <= rows() - row_offset);
    slp_assert(block_cols >= 0 && block_cols <= cols() - col_offset);
    return VariableBlock{*this, row_offset, col_offset, block_rows, block_cols};
  }
 
  const VariableBlock<const VariableMatrix> block(int row_offset,
                                                  int col_offset,
                                                  int block_rows,
                                                  int block_cols) const {
    slp_assert(row_offset >= 0 && row_offset <= rows());
    slp_assert(col_offset >= 0 && col_offset <= cols());
    slp_assert(block_rows >= 0 && block_rows <= rows() - row_offset);
    slp_assert(block_cols >= 0 && block_cols <= cols() - col_offset);
    return VariableBlock{*this, row_offset, col_offset, block_rows, block_cols};
  }
 
  VariableBlock<VariableMatrix> operator[](Slice row_slice, Slice col_slice) {
    int row_slice_length = row_slice.adjust(rows());
    int col_slice_length = col_slice.adjust(cols());
    return VariableBlock{*this, std::move(row_slice), row_slice_length,
                         std::move(col_slice), col_slice_length};
  }
 
  const VariableBlock<const VariableMatrix> operator[](Slice row_slice,
                                                       Slice col_slice) const {
    int row_slice_length = row_slice.adjust(rows());
    int col_slice_length = col_slice.adjust(cols());
    return VariableBlock{*this, std::move(row_slice), row_slice_length,
                         std::move(col_slice), col_slice_length};
  }
 
  VariableBlock<VariableMatrix> operator[](Slice row_slice,
                                           int row_slice_length,
                                           Slice col_slice,
                                           int col_slice_length) {
    return VariableBlock{*this, std::move(row_slice), row_slice_length,
                         std::move(col_slice), col_slice_length};
  }
 
  const VariableBlock<const VariableMatrix> operator[](
      Slice row_slice, int row_slice_length, Slice col_slice,
      int col_slice_length) const {
    return VariableBlock{*this, std::move(row_slice), row_slice_length,
                         std::move(col_slice), col_slice_length};
  }
 
  VariableBlock<VariableMatrix> segment(int offset, int length) {
    slp_assert(cols() == 1);
    slp_assert(offset >= 0 && offset < rows());
    slp_assert(length >= 0 && length <= rows() - offset);
    return block(offset, 0, length, 1);
  }
 
  const VariableBlock<const VariableMatrix> segment(int offset,
                                                    int length) const {
    slp_assert(cols() == 1);
    slp_assert(offset >= 0 && offset < rows());
    slp_assert(length >= 0 && length <= rows() - offset);
    return block(offset, 0, length, 1);
  }
 
  VariableBlock<VariableMatrix> row(int row) {
    slp_assert(row >= 0 && row < rows());
    return block(row, 0, 1, cols());
  }
 
  const VariableBlock<const VariableMatrix> row(int row) const {
    slp_assert(row >= 0 && row < rows());
    return block(row, 0, 1, cols());
  }
 
  VariableBlock<VariableMatrix> col(int col) {
    slp_assert(col >= 0 && col < cols());
    return block(0, col, rows(), 1);
  }
 
  const VariableBlock<const VariableMatrix> col(int col) const {
    slp_assert(col >= 0 && col < cols());
    return block(0, col, rows(), 1);
  }
 
  template <EigenMatrixLike LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.cols() == rhs.rows());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), rhs.cols());
 
    for (int i = 0; i < lhs.rows(); ++i) {
      for (int j = 0; j < rhs.cols(); ++j) {
        Variable sum{Scalar(0)};
        for (int k = 0; k < lhs.cols(); ++k) {
          sum += lhs[i, k] * rhs[k, j];
        }
        result[i, j] = sum;
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, EigenMatrixLike RHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.cols() == rhs.rows());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), rhs.cols());
 
    for (int i = 0; i < lhs.rows(); ++i) {
      for (int j = 0; j < rhs.cols(); ++j) {
        Variable sum{Scalar(0)};
        for (int k = 0; k < lhs.cols(); ++k) {
          sum += lhs[i, k] * rhs[k, j];
        }
        result[i, j] = sum;
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.cols() == rhs.rows());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), rhs.cols());
 
    for (int i = 0; i < lhs.rows(); ++i) {
      for (int j = 0; j < rhs.cols(); ++j) {
        Variable sum{Scalar(0)};
        for (int k = 0; k < lhs.cols(); ++k) {
          sum += lhs[i, k] * rhs[k, j];
        }
        result[i, j] = sum;
      }
    }
 
    return result;
  }
 
  template <EigenMatrixLike LHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs,
                                          const Variable<Scalar>& rhs) {
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] * rhs;
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, ScalarLike RHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs, const RHS& rhs) {
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] * rhs;
      }
    }
 
    return result;
  }
 
  template <EigenMatrixLike RHS>
  friend VariableMatrix<Scalar> operator*(const Variable<Scalar>& lhs,
                                          const RHS& rhs) {
    VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = rhs[row, col] * lhs;
      }
    }
 
    return result;
  }
 
  template <ScalarLike LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator*(const LHS& lhs, const RHS& rhs) {
    VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = rhs[row, col] * lhs;
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator*=(const MatrixLike auto& rhs) {
    slp_assert(cols() == rhs.rows() && cols() == rhs.cols());
 
    for (int i = 0; i < rows(); ++i) {
      VariableMatrix lhs_old_row = row(i);
      for (int j = 0; j < rhs.cols(); ++j) {
        Variable sum{Scalar(0)};
        for (int k = 0; k < cols(); ++k) {
          sum += lhs_old_row[k] * rhs[k, j];
        }
        (*this)[i, j] = sum;
      }
    }
 
    return *this;
  }
 
  VariableMatrix& operator*=(const ScalarLike auto& rhs) {
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] *= rhs;
      }
    }
 
    return *this;
  }
 
  template <EigenMatrixLike LHS>
  friend VariableMatrix<Scalar> operator/(const LHS& lhs,
                                          const Variable<Scalar>& rhs) {
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] / rhs;
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, ScalarLike RHS>
    requires(!SleipnirScalarLike<RHS, Scalar>)
  friend VariableMatrix<Scalar> operator/(const LHS& lhs, const RHS& rhs) {
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] / rhs;
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS>
  friend VariableMatrix<Scalar> operator/(const LHS& lhs,
                                          const Variable<Scalar>& rhs) {
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] / rhs;
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator/=(const ScalarLike auto& rhs) {
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] /= rhs;
      }
    }
 
    return *this;
  }
 
  template <EigenMatrixLike LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator+(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] + rhs[row, col];
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, EigenMatrixLike RHS>
  friend VariableMatrix<Scalar> operator+(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] + rhs[row, col];
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator+(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] + rhs[row, col];
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator+=(const MatrixLike auto& rhs) {
    slp_assert(rows() == rhs.rows() && cols() == rhs.cols());
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] += rhs[row, col];
      }
    }
 
    return *this;
  }
 
  VariableMatrix& operator+=(const ScalarLike auto& rhs) {
    slp_assert(rows() == 1 && cols() == 1);
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] += rhs;
      }
    }
 
    return *this;
  }
 
  template <EigenMatrixLike LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator-(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] - rhs[row, col];
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, EigenMatrixLike RHS>
  friend VariableMatrix<Scalar> operator-(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] - rhs[row, col];
      }
    }
 
    return result;
  }
 
  template <SleipnirMatrixLike<Scalar> LHS, SleipnirMatrixLike<Scalar> RHS>
  friend VariableMatrix<Scalar> operator-(const LHS& lhs, const RHS& rhs) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = lhs[row, col] - rhs[row, col];
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator-=(const MatrixLike auto& rhs) {
    slp_assert(rows() == rhs.rows() && cols() == rhs.cols());
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] -= rhs[row, col];
      }
    }
 
    return *this;
  }
 
  VariableMatrix& operator-=(const ScalarLike auto& rhs) {
    slp_assert(rows() == 1 && cols() == 1);
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        (*this)[row, col] -= rhs;
      }
    }
 
    return *this;
  }
 
  friend VariableMatrix<Scalar> operator-(
      const SleipnirMatrixLike<Scalar> auto& lhs) {
    VariableMatrix<Scalar> result{detail::empty, lhs.rows(), lhs.cols()};
 
    for (int row = 0; row < result.rows(); ++row) {
      for (int col = 0; col < result.cols(); ++col) {
        result[row, col] = -lhs[row, col];
      }
    }
 
    return result;
  }
 
  // NOLINTNEXTLINE (google-explicit-constructor)
  operator Variable<Scalar>() const {
    slp_assert(rows() == 1 && cols() == 1);
    return (*this)[0, 0];
  }
 
  VariableMatrix<Scalar> T() const {
    VariableMatrix<Scalar> result{detail::empty, cols(), rows()};
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        result[col, row] = (*this)[row, col];
      }
    }
 
    return result;
  }
 
  int rows() const { return m_rows; }
 
  int cols() const { return m_cols; }
 
  Scalar value(int row, int col) { return (*this)[row, col].value(); }
 
  Scalar value(int index) { return (*this)[index].value(); }
 
  Eigen::MatrixX<Scalar> value() {
    Eigen::MatrixX<Scalar> result{rows(), cols()};
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        result[row, col] = value(row, col);
      }
    }
 
    return result;
  }
 
  Eigen::MatrixX<Variable<Scalar>> to_eigen() const {
    Eigen::MatrixX<Variable<Scalar>> result{rows(), cols()};
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        result[row, col] = (*this)[row, col];
      }
    }
 
    return result;
  }
 
  VariableMatrix<Scalar> cwise_transform(
      function_ref<Variable<Scalar>(const Variable<Scalar>& x)> unary_op)
      const {
    VariableMatrix<Scalar> result{detail::empty, rows(), cols()};
 
    for (int row = 0; row < rows(); ++row) {
      for (int col = 0; col < cols(); ++col) {
        result[row, col] = unary_op((*this)[row, col]);
      }
    }
 
    return result;
  }
 
  VariableMatrix<Scalar> exp() const {
    slp_assert(rows() == cols());
 
    // Coefficients for (13, 13) Padé approximant of exp(A) are from the
    // following program:
    //
    // ```python
    // #!/usr/bin/env python
    //
    // import mpmath as mp
    //
    // # https://en.wikipedia.org/wiki/IEEE_754#Basic_and_interchange_formats
    // mp.mp.prec = 113  # quad precision
    //
    // L = 13
    // M = 13
    // p, q = mp.pade(mp.taylor(mp.exp, 0, L + M), L, M)
    //
    // print("constexpr std::array p{")
    // for k, p_k in enumerate(p):
    //     print(f"Scalar({p_k}L){',' if k < len(p) - 1 else '};'}")
    // print("constexpr std::array q{")
    // for k, q_k in enumerate(q):
    //     print(f"Scalar({q_k}L){',' if k < len(q) - 1 else '};'}")
    // ```
    constexpr size_t NUM_COEFFS = 14;
    constexpr std::array p{Scalar(1.0L),
                           Scalar(0.499999999999999999999987615564159L),
                           Scalar(0.119999999999999999999993719796468L),
                           Scalar(0.0183333333333333333333318078905823L),
                           Scalar(0.00199275362318840579710121408775205L),
                           Scalar(0.00016304347826086956521736560868541L),
                           Scalar(1.03519668737060041407846479512981e-05L),
                           Scalar(5.17598343685300207039205144279119e-07L),
                           Scalar(2.04315135665250081725989398668247e-08L),
                           Scalar(6.30602270571759511499920806570486e-10L),
                           Scalar(1.48377004840414002705850442238441e-11L),
                           Scalar(2.52915349159796595521307855480881e-13L),
                           Scalar(2.81017054621996217245857568431558e-15L),
                           Scalar(1.54404975067030888596595580141692e-17L)};
    constexpr std::array q{Scalar(1.0L),
                           Scalar(-0.500000000000000000000012384435841L),
                           Scalar(0.120000000000000000000006104232309L),
                           Scalar(-0.0183333333333333333333347707904729L),
                           Scalar(0.00199275362318840579710166350137732L),
                           Scalar(-0.000163043478260869565217413851346892L),
                           Scalar(1.03519668737060041407885187519331e-05L),
                           Scalar(-5.17598343685300207039443859549262e-07L),
                           Scalar(2.04315135665250081726103781506842e-08L),
                           Scalar(-6.30602270571759511500345035969254e-10L),
                           Scalar(1.48377004840414002705969744674248e-11L),
                           Scalar(-2.52915349159796595521550648156349e-13L),
                           Scalar(2.81017054621996217246180850207342e-15L),
                           Scalar(-1.54404975067030888596810607146175e-17L)};
    static_assert(p.size() == NUM_COEFFS);
    static_assert(q.size() == NUM_COEFFS);
 
    //     13
    // P = Σ pₖAᵏ
    //    k=0
    //
    //     13
    // Q = Σ qₖAᵏ
    //    k=0
    VariableMatrix<Scalar> P{
        Eigen::VectorX<Scalar>::Constant(rows(), p[0]).asDiagonal()};
    VariableMatrix<Scalar> Q{
        Eigen::VectorX<Scalar>::Constant(rows(), q[0]).asDiagonal()};
    auto A_pow = *this;
    size_t k = 1;
    while (true) {
      P += p[k] * A_pow;
      Q += q[k] * A_pow;
 
      ++k;
      if (k < NUM_COEFFS) {
        A_pow *= *this;
      } else {
        break;
      }
    }
 
    // https://mpmath.org/doc/current/calculus/approximation.html#mpmath.pade
    // defines the Padé approximant as exp(A)Q ≈ P, so:
    //
    //   exp(A) ≈ P / Q
    //   exp(A) ≈ (Qᵀ \ Pᵀ)ᵀ
    return VariableMatrix<Scalar>{
        Q.T().to_eigen().lu().solve(P.T().to_eigen()).transpose()};
  }
 
#ifndef DOXYGEN_SHOULD_SKIP_THIS
 
  class iterator {
   public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = Variable<Scalar>;
    using difference_type = std::ptrdiff_t;
    using pointer = Variable<Scalar>*;
    using reference = Variable<Scalar>&;
 
    constexpr iterator() noexcept = default;
 
    explicit constexpr iterator(
        gch::small_vector<Variable<Scalar>>::iterator it) noexcept
        : m_it{it} {}
 
    constexpr iterator& operator++() noexcept {
      ++m_it;
      return *this;
    }
 
    constexpr iterator operator++(int) noexcept {
      iterator retval = *this;
      ++(*this);
      return retval;
    }
 
    constexpr iterator& operator--() noexcept {
      --m_it;
      return *this;
    }
 
    constexpr iterator operator--(int) noexcept {
      iterator retval = *this;
      --(*this);
      return retval;
    }
 
    constexpr bool operator==(const iterator&) const noexcept = default;
 
    constexpr reference operator*() const noexcept { return *m_it; }
 
   private:
    gch::small_vector<Variable<Scalar>>::iterator m_it;
  };
 
  class const_iterator {
   public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = Variable<Scalar>;
    using difference_type = std::ptrdiff_t;
    using pointer = Variable<Scalar>*;
    using const_reference = const Variable<Scalar>&;
 
    constexpr const_iterator() noexcept = default;
 
    explicit constexpr const_iterator(
        gch::small_vector<Variable<Scalar>>::const_iterator it) noexcept
        : m_it{it} {}
 
    constexpr const_iterator& operator++() noexcept {
      ++m_it;
      return *this;
    }
 
    constexpr const_iterator operator++(int) noexcept {
      const_iterator retval = *this;
      ++(*this);
      return retval;
    }
 
    constexpr const_iterator& operator--() noexcept {
      --m_it;
      return *this;
    }
 
    constexpr const_iterator operator--(int) noexcept {
      const_iterator retval = *this;
      --(*this);
      return retval;
    }
 
    constexpr bool operator==(const const_iterator&) const noexcept = default;
 
    constexpr const_reference operator*() const noexcept { return *m_it; }
 
   private:
    gch::small_vector<Variable<Scalar>>::const_iterator m_it;
  };
 
  using reverse_iterator = std::reverse_iterator<iterator>;
  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
 
#endif  // DOXYGEN_SHOULD_SKIP_THIS
 
  iterator begin() { return iterator{m_storage.begin()}; }
 
  iterator end() { return iterator{m_storage.end()}; }
 
  const_iterator begin() const { return const_iterator{m_storage.begin()}; }
 
  const_iterator end() const { return const_iterator{m_storage.end()}; }
 
  const_iterator cbegin() const { return const_iterator{m_storage.cbegin()}; }
 
  const_iterator cend() const { return const_iterator{m_storage.cend()}; }
 
  reverse_iterator rbegin() { return reverse_iterator{end()}; }
 
  reverse_iterator rend() { return reverse_iterator{begin()}; }
 
  const_reverse_iterator rbegin() const {
    return const_reverse_iterator{end()};
  }
 
  const_reverse_iterator rend() const {
    return const_reverse_iterator{begin()};
  }
 
  const_reverse_iterator crbegin() const {
    return const_reverse_iterator{cend()};
  }
 
  const_reverse_iterator crend() const {
    return const_reverse_iterator{cbegin()};
  }
 
  size_t size() const { return m_storage.size(); }
 
  static VariableMatrix<Scalar> identity(int rows) {
    VariableMatrix<Scalar> result{detail::empty, rows, rows};
 
    for (int row = 0; row < rows; ++row) {
      for (int col = 0; col < rows; ++col) {
        result[row, col] = row == col ? Scalar(1) : Scalar(0);
      }
    }
 
    return result;
  }
 
  static VariableMatrix<Scalar> zero(int rows, int cols) {
    VariableMatrix<Scalar> result{detail::empty, rows, cols};
 
    for (auto& elem : result) {
      elem = Scalar(0);
    }
 
    return result;
  }
 
  static VariableMatrix<Scalar> one(int rows, int cols) {
    VariableMatrix<Scalar> result{detail::empty, rows, cols};
 
    for (auto& elem : result) {
      elem = Scalar(1);
    }
 
    return result;
  }
 
  static VariableMatrix<Scalar> constant(int rows, int cols, Scalar constant) {
    VariableMatrix<Scalar> result{detail::empty, rows, cols};
 
    for (auto& elem : result) {
      elem = constant;
    }
 
    return result;
  }
 
 private:
  gch::small_vector<Variable<Scalar>> m_storage;
  int m_rows = 0;
  int m_cols = 0;
};
 
template <typename Derived>
VariableMatrix(const Eigen::MatrixBase<Derived>&)
    -> VariableMatrix<typename Derived::Scalar>;
 
template <typename Derived>
VariableMatrix(const Eigen::DiagonalBase<Derived>&)
    -> VariableMatrix<typename Derived::Scalar>;
 
template <typename Scalar>
VariableMatrix<Scalar> cwise_reduce(
    const VariableMatrix<Scalar>& lhs, const VariableMatrix<Scalar>& rhs,
    function_ref<Variable<Scalar>(const Variable<Scalar>& x,
                                  const Variable<Scalar>& y)>
        binary_op) {
  slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
 
  VariableMatrix<Scalar> result{detail::empty, lhs.rows(), lhs.cols()};
 
  for (int row = 0; row < lhs.rows(); ++row) {
    for (int col = 0; col < lhs.cols(); ++col) {
      result[row, col] = binary_op(lhs[row, col], rhs[row, col]);
    }
  }
 
  return result;
}
 
template <typename Scalar>
VariableMatrix<Scalar> block(
    std::initializer_list<std::initializer_list<VariableMatrix<Scalar>>> list) {
  // Get row and column counts for destination matrix
  int rows = 0;
  int cols = -1;
  for (const auto& row : list) {
    if (row.size() > 0) {
      rows += row.begin()->rows();
    }
 
    // Get number of columns in this row
    int latest_cols = 0;
    for (const auto& elem : row) {
      // Assert the first and latest row have the same height
      slp_assert(row.begin()->rows() == elem.rows());
 
      latest_cols += elem.cols();
    }
 
    // If this is the first row, record the column count. Otherwise, assert the
    // first and latest column counts are the same.
    if (cols == -1) {
      cols = latest_cols;
    } else {
      slp_assert(cols == latest_cols);
    }
  }
 
  VariableMatrix<Scalar> result{detail::empty, rows, cols};
 
  int row_offset = 0;
  for (const auto& row : list) {
    int col_offset = 0;
    for (const auto& elem : row) {
      result.block(row_offset, col_offset, elem.rows(), elem.cols()) = elem;
      col_offset += elem.cols();
    }
    if (row.size() > 0) {
      row_offset += row.begin()->rows();
    }
  }
 
  return result;
}
 
template <typename Scalar>
VariableMatrix<Scalar> block(
    const std::vector<std::vector<VariableMatrix<Scalar>>>& list) {
  // Get row and column counts for destination matrix
  int rows = 0;
  int cols = -1;
  for (const auto& row : list) {
    if (row.size() > 0) {
      rows += row.begin()->rows();
    }
 
    // Get number of columns in this row
    int latest_cols = 0;
    for (const auto& elem : row) {
      // Assert the first and latest row have the same height
      slp_assert(row.begin()->rows() == elem.rows());
 
      latest_cols += elem.cols();
    }
 
    // If this is the first row, record the column count. Otherwise, assert the
    // first and latest column counts are the same.
    if (cols == -1) {
      cols = latest_cols;
    } else {
      slp_assert(cols == latest_cols);
    }
  }
 
  VariableMatrix<Scalar> result{detail::empty, rows, cols};
 
  int row_offset = 0;
  for (const auto& row : list) {
    int col_offset = 0;
    for (const auto& elem : row) {
      result.block(row_offset, col_offset, elem.rows(), elem.cols()) = elem;
      col_offset += elem.cols();
    }
    if (row.size() > 0) {
      row_offset += row.begin()->rows();
    }
  }
 
  return result;
}
 
template <typename Scalar>
VariableMatrix<Scalar> solve(const VariableMatrix<Scalar>& A,
                             const VariableMatrix<Scalar>& B) {
  // m x n * n x p = m x p
  slp_assert(A.rows() == B.rows());
 
  if (A.rows() == 1 && A.cols() == 1) {
    // Compute optimal inverse instead of using Eigen's general solver
    return B[0, 0] / A[0, 0];
  } else if (A.rows() == 2 && A.cols() == 2) {
    // Compute optimal inverse instead of using Eigen's general solver
    //
    // [a  b]⁻¹  ___1___ [ d  −b]
    // [c  d]  = ad − bc [−c   a]
 
    const auto& a = A[0, 0];
    const auto& b = A[0, 1];
    const auto& c = A[1, 0];
    const auto& d = A[1, 1];
 
    VariableMatrix adj_A{{d, -b}, {-c, a}};
    auto det_A = a * d - b * c;
    return adj_A / det_A * B;
  } else if (A.rows() == 3 && A.cols() == 3) {
    // Compute optimal inverse instead of using Eigen's general solver
    //
    // [a  b  c]⁻¹
    // [d  e  f]
    // [g  h  i]
    //                        1                   [ei − fh  ch − bi  bf − ce]
    //     = ------------------------------------ [fg − di  ai − cg  cd − af]
    //       a(ei − fh) + b(fg − di) + c(dh − eg) [dh − eg  bg − ah  ae − bd]
    //
    // https://www.wolframalpha.com/input?i=inverse+%7B%7Ba%2C+b%2C+c%7D%2C+%7Bd%2C+e%2C+f%7D%2C+%7Bg%2C+h%2C+i%7D%7D
 
    const auto& a = A[0, 0];
    const auto& b = A[0, 1];
    const auto& c = A[0, 2];
    const auto& d = A[1, 0];
    const auto& e = A[1, 1];
    const auto& f = A[1, 2];
    const auto& g = A[2, 0];
    const auto& h = A[2, 1];
    const auto& i = A[2, 2];
 
    auto ae = a * e;
    auto af = a * f;
    auto ah = a * h;
    auto ai = a * i;
    auto bd = b * d;
    auto bf = b * f;
    auto bg = b * g;
    auto bi = b * i;
    auto cd = c * d;
    auto ce = c * e;
    auto cg = c * g;
    auto ch = c * h;
    auto dh = d * h;
    auto di = d * i;
    auto eg = e * g;
    auto ei = e * i;
    auto fg = f * g;
    auto fh = f * h;
 
    auto adj_A00 = ei - fh;
    auto adj_A10 = fg - di;
    auto adj_A20 = dh - eg;
 
    VariableMatrix adj_A{{adj_A00, ch - bi, bf - ce},
                         {adj_A10, ai - cg, cd - af},
                         {adj_A20, bg - ah, ae - bd}};
    auto det_A = a * adj_A00 + b * adj_A10 + c * adj_A20;
    return adj_A / det_A * B;
  } else if (A.rows() == 4 && A.cols() == 4) {
    // Compute optimal inverse instead of using Eigen's general solver
    //
    // [a  b  c  d]⁻¹
    // [e  f  g  h]
    // [i  j  k  l]
    // [m  n  o  p]
    //
    // https://www.wolframalpha.com/input?i=inverse+%7B%7Ba%2C+b%2C+c%2C+d%7D%2C+%7Be%2C+f%2C+g%2C+h%7D%2C+%7Bi%2C+j%2C+k%2C+l%7D%2C+%7Bm%2C+n%2C+o%2C+p%7D%7D
 
    const auto& a = A[0, 0];
    const auto& b = A[0, 1];
    const auto& c = A[0, 2];
    const auto& d = A[0, 3];
    const auto& e = A[1, 0];
    const auto& f = A[1, 1];
    const auto& g = A[1, 2];
    const auto& h = A[1, 3];
    const auto& i = A[2, 0];
    const auto& j = A[2, 1];
    const auto& k = A[2, 2];
    const auto& l = A[2, 3];
    const auto& m = A[3, 0];
    const auto& n = A[3, 1];
    const auto& o = A[3, 2];
    const auto& p = A[3, 3];
 
    auto afk = a * f * k;
    auto afl = a * f * l;
    auto afo = a * f * o;
    auto afp = a * f * p;
    auto agj = a * g * j;
    auto agl = a * g * l;
    auto agn = a * g * n;
    auto agp = a * g * p;
    auto ahj = a * h * j;
    auto ahk = a * h * k;
    auto ahn = a * h * n;
    auto aho = a * h * o;
    auto ajo = a * j * o;
    auto ajp = a * j * p;
    auto akn = a * k * n;
    auto akp = a * k * p;
    auto aln = a * l * n;
    auto alo = a * l * o;
    auto bek = b * e * k;
    auto bel = b * e * l;
    auto beo = b * e * o;
    auto bep = b * e * p;
    auto bgi = b * g * i;
    auto bgl = b * g * l;
    auto bgm = b * g * m;
    auto bgp = b * g * p;
    auto bhi = b * h * i;
    auto bhk = b * h * k;
    auto bhm = b * h * m;
    auto bho = b * h * o;
    auto bio = b * i * o;
    auto bip = b * i * p;
    auto bjp = b * j * p;
    auto bkm = b * k * m;
    auto bkp = b * k * p;
    auto blm = b * l * m;
    auto blo = b * l * o;
    auto cej = c * e * j;
    auto cel = c * e * l;
    auto cen = c * e * n;
    auto cep = c * e * p;
    auto cfi = c * f * i;
    auto cfl = c * f * l;
    auto cfm = c * f * m;
    auto cfp = c * f * p;
    auto chi = c * h * i;
    auto chj = c * h * j;
    auto chm = c * h * m;
    auto chn = c * h * n;
    auto cin = c * i * n;
    auto cip = c * i * p;
    auto cjm = c * j * m;
    auto cjp = c * j * p;
    auto clm = c * l * m;
    auto cln = c * l * n;
    auto dej = d * e * j;
    auto dek = d * e * k;
    auto den = d * e * n;
    auto deo = d * e * o;
    auto dfi = d * f * i;
    auto dfk = d * f * k;
    auto dfm = d * f * m;
    auto dfo = d * f * o;
    auto dgi = d * g * i;
    auto dgj = d * g * j;
    auto dgm = d * g * m;
    auto dgn = d * g * n;
    auto din = d * i * n;
    auto dio = d * i * o;
    auto djm = d * j * m;
    auto djo = d * j * o;
    auto dkm = d * k * m;
    auto dkn = d * k * n;
    auto ejo = e * j * o;
    auto ejp = e * j * p;
    auto ekn = e * k * n;
    auto ekp = e * k * p;
    auto eln = e * l * n;
    auto elo = e * l * o;
    auto fio = f * i * o;
    auto fip = f * i * p;
    auto fkm = f * k * m;
    auto fkp = f * k * p;
    auto flm = f * l * m;
    auto flo = f * l * o;
    auto gin = g * i * n;
    auto gip = g * i * p;
    auto gjm = g * j * m;
    auto gjp = g * j * p;
    auto glm = g * l * m;
    auto gln = g * l * n;
    auto hin = h * i * n;
    auto hio = h * i * o;
    auto hjm = h * j * m;
    auto hjo = h * j * o;
    auto hkm = h * k * m;
    auto hkn = h * k * n;
 
    auto adj_A00 = fkp - flo - gjp + gln + hjo - hkn;
    auto adj_A01 = -bkp + blo + cjp - cln - djo + dkn;
    auto adj_A02 = bgp - bho - cfp + chn + dfo - dgn;
    auto adj_A03 = -bgl + bhk + cfl - chj - dfk + dgj;
    auto adj_A10 = -ekp + elo + gip - glm - hio + hkm;
    auto adj_A11 = akp - alo - cip + clm + dio - dkm;
    auto adj_A12 = -agp + aho + cep - chm - deo + dgm;
    auto adj_A13 = agl - ahk - cel + chi + dek - dgi;
    auto adj_A20 = ejp - eln - fip + flm + hin - hjm;
    auto adj_A21 = -ajp + aln + bip - blm - din + djm;
    auto adj_A22 = afp - ahn - bep + bhm + den - dfm;
    auto adj_A23 = -afl + ahj + bel - bhi - dej + dfi;
    auto adj_A30 = -ejo + ekn + fio - fkm - gin + gjm;
    // NOLINTNEXTLINE(build/include_what_you_use)
    auto adj_A31 = ajo - akn - bio + bkm + cin - cjm;
    auto adj_A32 = -afo + agn + beo - bgm - cen + cfm;
    auto adj_A33 = afk - agj - bek + bgi + cej - cfi;
 
    VariableMatrix adj_A{{adj_A00, adj_A01, adj_A02, adj_A03},
                         {adj_A10, adj_A11, adj_A12, adj_A13},
                         {adj_A20, adj_A21, adj_A22, adj_A23},
                         {adj_A30, adj_A31, adj_A32, adj_A33}};
    auto det_A = a * adj_A00 + b * adj_A10 + c * adj_A20 + d * adj_A30;
    return adj_A / det_A * B;
  } else {
    return VariableMatrix<Scalar>{
        A.to_eigen().householderQr().solve(B.to_eigen())};
  }
}
 
namespace detail {
 
template <typename Scalar>
VariableMatrix<Scalar> gradient_tree(const ExpressionGraph<Scalar>& top_list,
                                     const VariableMatrix<Scalar>& wrt) {
  slp_assert(wrt.cols() == 1);
 
  // Read docs/algorithms.md#Reverse_accumulation_automatic_differentiation
  // for background on reverse accumulation automatic differentiation.
 
  if (top_list.empty()) {
    return VariableMatrix<Scalar>{detail::empty, wrt.rows(), 1};
  }
 
  // Set root node's adjoint to 1 since df/df is 1
  top_list[0]->adj_expr = constant_ptr(Scalar(1));
 
  // df/dx = (df/dy)(dy/dx). The adjoint of x is equal to the adjoint of y
  // multiplied by dy/dx. If there are multiple "paths" from the root node to
  // variable; the variable's adjoint is the sum of each path's adjoint
  // contribution.
  for (auto& node : top_list) {
    // Leaf nodes have no children to propagate adjoints to, and a null (zero)
    // adjoint has nothing to contribute
    if (!node->is_leaf && node->adj_expr != nullptr) {
      node->accumulate_adjoints_expr();
    }
  }
 
  // Move gradient tree to return value
  VariableMatrix<Scalar> grad{detail::empty, wrt.rows(), 1};
  for (int row = 0; row < grad.rows(); ++row) {
    grad[row] = Variable{std::move(wrt[row].expr->adj_expr)};
  }
 
  // Unlink adjoints to avoid circular references between them and their
  // parent expressions. This ensures all expressions are returned to the free
  // list.
  for (auto& node : top_list) {
    node->adj_expr = nullptr;
  }
 
  return grad;
}
 
}  // namespace detail
 
extern template class EXPORT_TEMPLATE_DECLARE(SLEIPNIR_DLLEXPORT)
VariableMatrix<double>;
 
extern template SLEIPNIR_DLLEXPORT VariableMatrix<double> solve(
    const VariableMatrix<double>& A, const VariableMatrix<double>& B);
 
}  // namespace slp