expression.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <stdint.h>
 
#include <algorithm>
#include <cmath>
#include <memory>
#include <numbers>
#include <string_view>
#include <type_traits>
#include <utility>
 
#include <gch/small_vector.hpp>
 
#include "sleipnir/autodiff/expression_type.hpp"
#include "sleipnir/util/function_ref.hpp"
#include "sleipnir/util/intrusive_shared_ptr.hpp"
#include "sleipnir/util/pool.hpp"
 
namespace slp::detail {
 
#ifdef _WIN32
inline constexpr bool USE_POOL_ALLOCATOR = false;
#else
inline constexpr bool USE_POOL_ALLOCATOR = true;
#endif
 
template <typename Scalar>
struct Expression;
 
template <typename Scalar>
using ExpressionPtr = IntrusiveSharedPtr<Expression<Scalar>>;
 
template <typename T, typename... Args>
static ExpressionPtr<typename T::Scalar> make_expression_ptr(Args&&... args) {
  if constexpr (USE_POOL_ALLOCATOR) {
    return allocate_intrusive_shared<T>(global_pool_allocator<T>(),
                                        std::forward<Args>(args)...);
  } else {
    return make_intrusive_shared<T>(std::forward<Args>(args)...);
  }
}
 
template <typename Scalar, ExpressionType T>
struct BinaryMinusExpression;
 
template <typename Scalar, ExpressionType T>
struct BinaryPlusExpression;
 
template <typename Scalar, ExpressionType T>
struct DivExpression;
 
template <typename Scalar, ExpressionType T>
struct MultExpression;
 
template <typename Scalar, ExpressionType T>
struct UnaryMinusExpression;
 
template <typename Scalar>
ExpressionPtr<Scalar> constant_ptr(Scalar value);
 
template <typename Scalar_>
struct Expression {
  using Scalar = Scalar_;
 
  Scalar val{0};
 
  Scalar adj{0};
 
  ExpressionPtr<Scalar> adj_expr;
 
  bool is_leaf = false;
 
  int32_t scratch = -1;
 
  uint32_t ref_count = 0;
 
  constexpr Expression() = default;
 
  explicit constexpr Expression(Scalar value) : val{value}, is_leaf{true} {}
 
  virtual ~Expression() = default;
 
  constexpr bool is_constant(Scalar constant) const {
    return type() == ExpressionType::CONSTANT && val == constant;
  }
 
  friend ExpressionPtr<Scalar> operator*(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    } else if (rhs->is_constant(Scalar(0))) {
      // Return zero, which rhs currently is
      return rhs;
    } else if (lhs->is_constant(Scalar(1))) {
      // Return rhs unmodified
      return rhs;
    } else if (rhs->is_constant(Scalar(1))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val * rhs->val);
    }
 
    // Evaluate expression type
    if (lhs->type() == CONSTANT) {
      if (rhs->type() == LINEAR) {
        return make_expression_ptr<MultExpression<Scalar, LINEAR>>(lhs, rhs);
      } else if (rhs->type() == QUADRATIC) {
        return make_expression_ptr<MultExpression<Scalar, QUADRATIC>>(lhs, rhs);
      } else {
        return make_expression_ptr<MultExpression<Scalar, NONLINEAR>>(lhs, rhs);
      }
    } else if (rhs->type() == CONSTANT) {
      if (lhs->type() == LINEAR) {
        return make_expression_ptr<MultExpression<Scalar, LINEAR>>(lhs, rhs);
      } else if (lhs->type() == QUADRATIC) {
        return make_expression_ptr<MultExpression<Scalar, QUADRATIC>>(lhs, rhs);
      } else {
        return make_expression_ptr<MultExpression<Scalar, NONLINEAR>>(lhs, rhs);
      }
    } else if (lhs->type() == LINEAR && rhs->type() == LINEAR) {
      return make_expression_ptr<MultExpression<Scalar, QUADRATIC>>(lhs, rhs);
    } else {
      return make_expression_ptr<MultExpression<Scalar, NONLINEAR>>(lhs, rhs);
    }
  }
 
  friend ExpressionPtr<Scalar> operator/(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    } else if (rhs->is_constant(Scalar(1))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val / rhs->val);
    }
 
    // Evaluate expression type
    if (rhs->type() == CONSTANT) {
      if (lhs->type() == LINEAR) {
        return make_expression_ptr<DivExpression<Scalar, LINEAR>>(lhs, rhs);
      } else if (lhs->type() == QUADRATIC) {
        return make_expression_ptr<DivExpression<Scalar, QUADRATIC>>(lhs, rhs);
      } else {
        return make_expression_ptr<DivExpression<Scalar, NONLINEAR>>(lhs, rhs);
      }
    } else {
      return make_expression_ptr<DivExpression<Scalar, NONLINEAR>>(lhs, rhs);
    }
  }
 
  friend ExpressionPtr<Scalar> operator+(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression. We check for nullptr because operator+ is used in
    // adjoint accumulation, and child nodes can be null.
    if (lhs == nullptr || lhs->is_constant(Scalar(0))) {
      // Return rhs unmodified
      return rhs;
    } else if (rhs == nullptr || rhs->is_constant(Scalar(0))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val + rhs->val);
    }
 
    auto type = std::max(lhs->type(), rhs->type());
    if (type == LINEAR) {
      return make_expression_ptr<BinaryPlusExpression<Scalar, LINEAR>>(lhs,
                                                                       rhs);
    } else if (type == QUADRATIC) {
      return make_expression_ptr<BinaryPlusExpression<Scalar, QUADRATIC>>(lhs,
                                                                          rhs);
    } else {
      return make_expression_ptr<BinaryPlusExpression<Scalar, NONLINEAR>>(lhs,
                                                                          rhs);
    }
  }
 
  friend ExpressionPtr<Scalar> operator+=(ExpressionPtr<Scalar>& lhs,
                                          const ExpressionPtr<Scalar>& rhs) {
    return lhs = lhs + rhs;
  }
 
  friend ExpressionPtr<Scalar> operator-(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      if (rhs->is_constant(Scalar(0))) {
        // Return zero, which rhs currently is
        return rhs;
      } else {
        // Return rhs negated
        return -rhs;
      }
    } else if (rhs->is_constant(Scalar(0))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val - rhs->val);
    }
 
    auto type = std::max(lhs->type(), rhs->type());
    if (type == LINEAR) {
      return make_expression_ptr<BinaryMinusExpression<Scalar, LINEAR>>(lhs,
                                                                        rhs);
    } else if (type == QUADRATIC) {
      return make_expression_ptr<BinaryMinusExpression<Scalar, QUADRATIC>>(lhs,
                                                                           rhs);
    } else {
      return make_expression_ptr<BinaryMinusExpression<Scalar, NONLINEAR>>(lhs,
                                                                           rhs);
    }
  }
 
  friend ExpressionPtr<Scalar> operator-(const ExpressionPtr<Scalar>& lhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT) {
      return constant_ptr(-lhs->val);
    }
 
    if (lhs->type() == LINEAR) {
      return make_expression_ptr<UnaryMinusExpression<Scalar, LINEAR>>(lhs);
    } else if (lhs->type() == QUADRATIC) {
      return make_expression_ptr<UnaryMinusExpression<Scalar, QUADRATIC>>(lhs);
    } else {
      return make_expression_ptr<UnaryMinusExpression<Scalar, NONLINEAR>>(lhs);
    }
  }
 
  friend ExpressionPtr<Scalar> operator+(const ExpressionPtr<Scalar>& lhs) {
    return lhs;
  }
 
  virtual void visit_args(
      [[maybe_unused]] function_ref<void(Expression<Scalar>* arg)> func) const {
  }
 
  virtual Scalar value() const = 0;
 
  virtual ExpressionType type() const = 0;
 
  virtual std::string_view name() const = 0;
 
  virtual void accumulate_adjoints() const {}
 
  virtual void accumulate_adjoints_expr() const {}
};
 
template <typename Scalar, ExpressionType T>
struct BinaryMinusExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> lhs;
 
  ExpressionPtr<Scalar> rhs;
 
  constexpr BinaryMinusExpression(ExpressionPtr<Scalar> lhs,
                                  ExpressionPtr<Scalar> rhs)
      : lhs{std::move(lhs)}, rhs{std::move(rhs)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(lhs.get());
    func(rhs.get());
  }
 
  Scalar value() const override { return lhs->val - rhs->val; }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "binary minus"; }
 
  void accumulate_adjoints() const override {
    lhs->adj += grad_l();
    rhs->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    lhs->adj_expr += grad_expr_l();
    rhs->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return this->adj; }
 
  Scalar grad_r() const { return -this->adj; }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr; }
 
  ExpressionPtr<Scalar> grad_expr_r() const { return -this->adj_expr; }
};
 
template <typename Scalar, ExpressionType T>
struct BinaryPlusExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> lhs;
 
  ExpressionPtr<Scalar> rhs;
 
  constexpr BinaryPlusExpression(ExpressionPtr<Scalar> lhs,
                                 ExpressionPtr<Scalar> rhs)
      : lhs{std::move(lhs)}, rhs{std::move(rhs)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(lhs.get());
    func(rhs.get());
  }
 
  Scalar value() const override { return lhs->val + rhs->val; }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "binary plus"; }
 
  void accumulate_adjoints() const override {
    lhs->adj += grad_l();
    rhs->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    lhs->adj_expr += grad_expr_l();
    rhs->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return this->adj; }
 
  Scalar grad_r() const { return this->adj; }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr; }
 
  ExpressionPtr<Scalar> grad_expr_r() const { return this->adj_expr; }
};
 
template <typename Scalar>
struct CbrtExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr CbrtExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::cbrt;
    return cbrt(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "cbrt"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::cbrt;
 
    Scalar c = cbrt(x->val);
    return this->adj / (Scalar(3) * c * c);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    auto c = cbrt(x);
    return this->adj_expr / (constant_ptr(Scalar(3)) * c * c);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cbrt(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::cbrt;
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    if (x->val == Scalar(0)) {
      // Return zero
      return x;
    } else if (x->val == Scalar(-1) || x->val == Scalar(1)) {
      return x;
    } else {
      return constant_ptr(cbrt(x->val));
    }
  }
 
  return make_expression_ptr<CbrtExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct ConstantExpression final : Expression<Scalar> {
  explicit constexpr ConstantExpression(Scalar value)
      : Expression<Scalar>{value} {}
 
  Scalar value() const override { return this->val; }
 
  ExpressionType type() const override { return ExpressionType::CONSTANT; }
 
  std::string_view name() const override { return "constant"; }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> constant_ptr(Scalar value) {
  return make_expression_ptr<ConstantExpression<Scalar>>(value);
}
 
template <typename Scalar>
struct DecisionVariableExpression final : Expression<Scalar> {
  constexpr DecisionVariableExpression() : Expression<Scalar>{Scalar(0)} {}
 
  explicit constexpr DecisionVariableExpression(Scalar value)
      : Expression<Scalar>{value} {}
 
  Scalar value() const override { return this->val; }
 
  ExpressionType type() const override { return ExpressionType::LINEAR; }
 
  std::string_view name() const override { return "decision variable"; }
};
 
template <typename Scalar, ExpressionType T>
struct DivExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> lhs;
 
  ExpressionPtr<Scalar> rhs;
 
  constexpr DivExpression(ExpressionPtr<Scalar> lhs, ExpressionPtr<Scalar> rhs)
      : lhs{std::move(lhs)}, rhs{std::move(rhs)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(lhs.get());
    func(rhs.get());
  }
 
  Scalar value() const override { return lhs->val / rhs->val; }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "division"; }
 
  void accumulate_adjoints() const override {
    lhs->adj += grad_l();
    rhs->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    lhs->adj_expr += grad_expr_l();
    rhs->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return this->adj / rhs->val; };
 
  Scalar grad_r() const {
    return this->adj * -lhs->val / (rhs->val * rhs->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr / rhs; }
 
  ExpressionPtr<Scalar> grad_expr_r() const {
    return this->adj_expr * -lhs / (rhs * rhs);
  }
};
 
template <typename Scalar, ExpressionType T>
struct MultExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> lhs;
 
  ExpressionPtr<Scalar> rhs;
 
  constexpr MultExpression(ExpressionPtr<Scalar> lhs, ExpressionPtr<Scalar> rhs)
      : lhs{std::move(lhs)}, rhs{std::move(rhs)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(lhs.get());
    func(rhs.get());
  }
 
  Scalar value() const override { return lhs->val * rhs->val; }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "multiplication"; }
 
  void accumulate_adjoints() const override {
    lhs->adj += grad_l();
    rhs->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    lhs->adj_expr += grad_expr_l();
    rhs->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return this->adj * rhs->val; }
 
  Scalar grad_r() const { return this->adj * lhs->val; }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * rhs; }
 
  ExpressionPtr<Scalar> grad_expr_r() const { return this->adj_expr * lhs; }
};
 
template <typename Scalar, ExpressionType T>
struct UnaryMinusExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> lhs;
 
  explicit constexpr UnaryMinusExpression(ExpressionPtr<Scalar> lhs)
      : lhs{std::move(lhs)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(lhs.get());
  }
 
  Scalar value() const override { return -lhs->val; }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "unary minus"; }
 
  void accumulate_adjoints() const override { lhs->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    lhs->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const { return -this->adj; }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return -this->adj_expr; }
};
 
template <typename Scalar>
constexpr void inc_ref_count(Expression<Scalar>* expr) {
  ++expr->ref_count;
}
 
template <typename Scalar>
constexpr void dec_ref_count(Expression<Scalar>* expr) {
  // If a deeply nested tree is being deallocated all at once, calling the
  // Expression destructor when expr's refcount reaches zero can cause a stack
  // overflow. Instead, we iterate over its children to decrement their
  // refcounts and deallocate them.
  gch::small_vector<Expression<Scalar>*> stack;
  stack.emplace_back(expr);
 
  while (!stack.empty()) {
    auto elem = stack.back();
    stack.pop_back();
 
    // Decrement the current node's refcount. If it reaches zero, deallocate the
    // node and enqueue its children so their refcounts are decremented too.
    if (--elem->ref_count == 0) {
      if (elem->adj_expr != nullptr) {
        stack.emplace_back(elem->adj_expr.get());
      }
      elem->visit_args([&stack](const auto& arg) { stack.emplace_back(arg); });
 
      // Not calling the destructor here is safe because it only decrements
      // refcounts, which was already done above.
      if constexpr (USE_POOL_ALLOCATOR) {
        auto alloc = global_pool_allocator<Expression<Scalar>>();
        std::allocator_traits<decltype(alloc)>::deallocate(alloc, elem, 1);
      } else {
        operator delete(elem);
      }
    }
  }
}
 
template <typename Scalar>
struct AbsExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr AbsExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::abs;
    return abs(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "abs"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    if (x->val < Scalar(0)) {
      return -this->adj;
    } else if (x->val > Scalar(0)) {
      return this->adj;
    } else {
      return Scalar(0);
    }
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * sign(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> abs(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::abs;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(abs(x->val));
  }
 
  return make_expression_ptr<AbsExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AcosExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr AcosExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::acos;
    return acos(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "acos"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::sqrt;
    return -this->adj / sqrt(Scalar(1) - x->val * x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return -this->adj_expr / sqrt(constant_ptr(Scalar(1)) - x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> acos(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::acos;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(std::numbers::pi) / Scalar(2));
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(acos(x->val));
  }
 
  return make_expression_ptr<AcosExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AsinExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr AsinExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::asin;
    return asin(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "asin"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::sqrt;
    return this->adj / sqrt(Scalar(1) - x->val * x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr / sqrt(constant_ptr(Scalar(1)) - x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> asin(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::asin;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(asin(x->val));
  }
 
  return make_expression_ptr<AsinExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AtanExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr AtanExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::atan;
    return atan(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "atan"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const { return this->adj / (Scalar(1) + x->val * x->val); }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr / (constant_ptr(Scalar(1)) + x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> atan(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::atan;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(atan(x->val));
  }
 
  return make_expression_ptr<AtanExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct Atan2Expression final : Expression<Scalar> {
  ExpressionPtr<Scalar> y;
 
  ExpressionPtr<Scalar> x;
 
  constexpr Atan2Expression(ExpressionPtr<Scalar> y, ExpressionPtr<Scalar> x)
      : y{std::move(y)}, x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(y.get());
    func(x.get());
  }
 
  Scalar value() const override {
    using std::atan2;
    return atan2(y->val, x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "atan2"; }
 
  void accumulate_adjoints() const override {
    y->adj += grad_l();
    x->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    y->adj_expr += grad_expr_l();
    x->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const {
    return this->adj * x->val / (y->val * y->val + x->val * x->val);
  }
 
  Scalar grad_r() const {
    return this->adj * -y->val / (y->val * y->val + x->val * x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr * x / (y * y + x * x);
  }
 
  ExpressionPtr<Scalar> grad_expr_r() const {
    return this->adj_expr * -y / (y * y + x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> atan2(const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::atan2;
 
  // Evaluate constant
  if (y->type() == CONSTANT && x->type() == CONSTANT) {
    return constant_ptr(atan2(y->val, x->val));
  }
 
  return make_expression_ptr<Atan2Expression<Scalar>>(y, x);
}
 
template <typename Scalar>
struct CosExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr CosExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::cos;
    return cos(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "cos"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::sin;
    return this->adj * -sin(x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * -sin(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cos(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::cos;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(1));
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(cos(x->val));
  }
 
  return make_expression_ptr<CosExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct CoshExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr CoshExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::cosh;
    return cosh(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "cosh"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::sinh;
    return this->adj * sinh(x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * sinh(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cosh(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::cosh;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(1));
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(cosh(x->val));
  }
 
  return make_expression_ptr<CoshExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct ErfExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr ErfExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::erf;
    return erf(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "erf"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::exp;
    return this->adj * Scalar(2.0 * std::numbers::inv_sqrtpi) *
           exp(-x->val * x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr *
           constant_ptr(Scalar(2.0 * std::numbers::inv_sqrtpi)) * exp(-x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> erf(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::erf;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(erf(x->val));
  }
 
  return make_expression_ptr<ErfExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct ExpExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr ExpExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::exp;
    return exp(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "exp"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::exp;
    return this->adj * exp(x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * exp(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> exp(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::exp;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(1));
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(exp(x->val));
  }
 
  return make_expression_ptr<ExpExpression<Scalar>>(x);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y);
 
template <typename Scalar>
struct Hypot2Expression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  ExpressionPtr<Scalar> y;
 
  constexpr Hypot2Expression(ExpressionPtr<Scalar> x, ExpressionPtr<Scalar> y)
      : x{std::move(x)}, y{std::move(y)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
    func(y.get());
  }
 
  Scalar value() const override {
    using std::hypot;
    return hypot(x->val, y->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "hypot"; }
 
  void accumulate_adjoints() const override {
    using std::hypot;
    Scalar norm = hypot(x->val, y->val);
    x->adj += this->adj * x->val / norm;
    y->adj += this->adj * y->val / norm;
  }
 
  void accumulate_adjoints_expr() const override {
    auto norm = hypot(x, y);
    x->adj_expr += this->adj_expr * x / norm;
    y->adj_expr += this->adj_expr * y / norm;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y) {
  using enum ExpressionType;
  using std::hypot;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return abs(y);
  } else if (y->is_constant(Scalar(0))) {
    return abs(x);
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT && y->type() == CONSTANT) {
    return constant_ptr(hypot(x->val, y->val));
  }
 
  return make_expression_ptr<Hypot2Expression<Scalar>>(x, y);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& z);
 
template <typename Scalar>
struct Hypot3Expression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  ExpressionPtr<Scalar> y;
 
  ExpressionPtr<Scalar> z;
 
  constexpr Hypot3Expression(ExpressionPtr<Scalar> x, ExpressionPtr<Scalar> y,
                             ExpressionPtr<Scalar> z)
      : x{std::move(x)}, y{std::move(y)}, z{std::move(z)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
    func(y.get());
    func(z.get());
  }
 
  Scalar value() const override {
    using std::hypot;
    return hypot(x->val, y->val, z->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "hypot"; }
 
  void accumulate_adjoints() const override {
    using std::hypot;
    Scalar norm = hypot(x->val, y->val, z->val);
    x->adj += this->adj * x->val / norm;
    y->adj += this->adj * y->val / norm;
    z->adj += this->adj * z->val / norm;
  }
 
  void accumulate_adjoints_expr() const override {
    auto norm = hypot(x, y, z);
    x->adj_expr += this->adj_expr * x / norm;
    y->adj_expr += this->adj_expr * y / norm;
    z->adj_expr += this->adj_expr * z / norm;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& z) {
  using enum ExpressionType;
  using std::hypot;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return hypot(y, z);
  } else if (y->is_constant(Scalar(0))) {
    return hypot(x, z);
  } else if (z->is_constant(Scalar(0))) {
    return hypot(x, y);
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT && y->type() == CONSTANT && z->type() == CONSTANT) {
    return constant_ptr(hypot(x->val, y->val, z->val));
  }
 
  return make_expression_ptr<Hypot3Expression<Scalar>>(x, y, z);
}
 
template <typename Scalar>
struct IfElseExpression final : Expression<Scalar> {
  bool (*cond)(Scalar a, Scalar b);
 
  ExpressionPtr<Scalar> a;
 
  ExpressionPtr<Scalar> b;
 
  ExpressionPtr<Scalar> t;
 
  ExpressionPtr<Scalar> f;
 
  constexpr IfElseExpression(bool (*cond)(Scalar a, Scalar b),
                             ExpressionPtr<Scalar> a, ExpressionPtr<Scalar> b,
                             ExpressionPtr<Scalar> t, ExpressionPtr<Scalar> f)
      : cond{cond},
        a{std::move(a)},
        b{std::move(b)},
        t{std::move(t)},
        f{std::move(f)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(a.get());
    func(b.get());
    func(t.get());
    func(f.get());
  }
 
  Scalar value() const override {
    return cond(a->val, b->val) ? t->val : f->val;
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "if-else"; }
 
  // a and b have zero gradient, so only t and f accumulate adjoints
  void accumulate_adjoints() const override {
    if (cond(a->val, b->val)) {
      t->adj += this->adj;
    } else {
      f->adj += this->adj;
    }
  }
 
  void accumulate_adjoints_expr() const override {
    t->adj_expr += grad_expr_t();
    f->adj_expr += grad_expr_f();
  }
 
 private:
  ExpressionPtr<Scalar> grad_expr_t() const {
    // adjoint if cond(a, b), otherwise 0
    return if_else(cond, a, b, this->adj_expr, constant_ptr(Scalar(0)));
  }
 
  ExpressionPtr<Scalar> grad_expr_f() const {
    // 0 if cond(a, b), otherwise adjoint
    return if_else(cond, a, b, constant_ptr(Scalar(0)), this->adj_expr);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> if_else(
    std::type_identity_t<bool (*)(Scalar a, Scalar b)> cond,
    const ExpressionPtr<Scalar>& a, const ExpressionPtr<Scalar>& b,
    const ExpressionPtr<Scalar>& t, const ExpressionPtr<Scalar>& f) {
  using enum ExpressionType;
 
  // Prune expression
  if (t == f) {
    // Return t, which both branches currently are
    return t;
  } else if (t->type() == CONSTANT && f->type() == CONSTANT &&
             t->val == f->val) {
    // Return t, which both branches currently equal
    return t;
  }
 
  // Evaluate constant condition
  if (a->type() == CONSTANT && b->type() == CONSTANT) {
    return cond(a->val, b->val) ? t : f;
  }
 
  return make_expression_ptr<IfElseExpression<Scalar>>(cond, a, b, t, f);
}
 
template <typename Scalar>
struct LogExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr LogExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::log;
    return log(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "log"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const { return this->adj / x->val; }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr / x; }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> log(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::log;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(log(x->val));
  }
 
  return make_expression_ptr<LogExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct Log10Expression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr Log10Expression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::log10;
    return log10(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "log10"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    return this->adj / (Scalar(std::numbers::ln10) * x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr / (constant_ptr(Scalar(std::numbers::ln10)) * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> log10(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::log10;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(log10(x->val));
  }
 
  return make_expression_ptr<Log10Expression<Scalar>>(x);
}
 
template <typename Scalar>
struct MaxExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> a;
 
  ExpressionPtr<Scalar> b;
 
  constexpr MaxExpression(ExpressionPtr<Scalar> a, ExpressionPtr<Scalar> b)
      : a{std::move(a)}, b{std::move(b)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(a.get());
    func(b.get());
  }
 
  Scalar value() const override {
    using std::max;
    return max(a->val, b->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "max"; }
 
  void accumulate_adjoints() const override {
    a->adj += grad_l();
    b->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    a->adj_expr += grad_expr_l();
    b->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return a->val >= b->val ? this->adj : Scalar(0); }
 
  Scalar grad_r() const { return a->val >= b->val ? Scalar(0) : this->adj; }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    // adjoint if a >= b, otherwise 0
    return if_else([](Scalar a, Scalar b) { return a >= b; }, a, b,
                   this->adj_expr, constant_ptr(Scalar(0)));
  }
 
  ExpressionPtr<Scalar> grad_expr_r() const {
    // 0 if a >= b, otherwise adjoint
    return if_else([](Scalar a, Scalar b) { return a >= b; }, a, b,
                   constant_ptr(Scalar(0)), this->adj_expr);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> max(const ExpressionPtr<Scalar>& a,
                          const ExpressionPtr<Scalar>& b) {
  using enum ExpressionType;
  using std::max;
 
  // Evaluate constant
  if (a->type() == CONSTANT && b->type() == CONSTANT) {
    return constant_ptr(max(a->val, b->val));
  }
 
  return make_expression_ptr<MaxExpression<Scalar>>(a, b);
}
 
template <typename Scalar>
struct MinExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> a;
 
  ExpressionPtr<Scalar> b;
 
  constexpr MinExpression(ExpressionPtr<Scalar> a, ExpressionPtr<Scalar> b)
      : a{std::move(a)}, b{std::move(b)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(a.get());
    func(b.get());
  }
 
  Scalar value() const override {
    using std::min;
    return min(a->val, b->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "min"; }
 
  void accumulate_adjoints() const override {
    a->adj += grad_l();
    b->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    a->adj_expr += grad_expr_l();
    b->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const { return a->val <= b->val ? this->adj : Scalar(0); }
 
  Scalar grad_r() const { return a->val <= b->val ? Scalar(0) : this->adj; }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    // adjoint if a <= b, otherwise 0
    return if_else([](Scalar a, Scalar b) { return a <= b; }, a, b,
                   this->adj_expr, constant_ptr(Scalar(0)));
  }
 
  ExpressionPtr<Scalar> grad_expr_r() const {
    // 0 if a <= b, otherwise adjoint
    return if_else([](Scalar a, Scalar b) { return a <= b; }, a, b,
                   constant_ptr(Scalar(0)), this->adj_expr);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> min(const ExpressionPtr<Scalar>& a,
                          const ExpressionPtr<Scalar>& b) {
  using enum ExpressionType;
  using std::min;
 
  // Evaluate constant
  if (a->type() == CONSTANT && b->type() == CONSTANT) {
    return constant_ptr(min(a->val, b->val));
  }
 
  return make_expression_ptr<MinExpression<Scalar>>(a, b);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> pow(const ExpressionPtr<Scalar>& base,
                          const ExpressionPtr<Scalar>& power);
 
template <typename Scalar, ExpressionType T>
struct PowExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> base;
 
  ExpressionPtr<Scalar> power;
 
  constexpr PowExpression(ExpressionPtr<Scalar> base,
                          ExpressionPtr<Scalar> power)
      : base{std::move(base)}, power{std::move(power)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(base.get());
    func(power.get());
  }
 
  Scalar value() const override {
    using std::pow;
    return pow(base->val, power->val);
  }
 
  ExpressionType type() const override { return T; }
 
  std::string_view name() const override { return "pow"; }
 
  void accumulate_adjoints() const override {
    base->adj += grad_l();
    power->adj += grad_r();
  }
 
  void accumulate_adjoints_expr() const override {
    base->adj_expr += grad_expr_l();
    power->adj_expr += grad_expr_r();
  }
 
 private:
  Scalar grad_l() const {
    using std::pow;
    return this->adj * pow(base->val, power->val - Scalar(1)) * power->val;
  }
 
  Scalar grad_r() const {
    using std::log;
    using std::pow;
 
    return this->adj * pow(base->val, power->val) * log(base->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr * pow(base, power - constant_ptr(Scalar(1))) * power;
  }
 
  ExpressionPtr<Scalar> grad_expr_r() const {
    return this->adj_expr * pow(base, power) * log(base);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> pow(const ExpressionPtr<Scalar>& base,
                          const ExpressionPtr<Scalar>& power) {
  using enum ExpressionType;
  using std::pow;
 
  // Prune expression
  if (base->is_constant(Scalar(0))) {
    // Return zero, which base currently is
    return base;
  } else if (base->is_constant(Scalar(1))) {
    // Return one, which base currently is
    return base;
  }
  if (power->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(1));
  } else if (power->is_constant(Scalar(1))) {
    // Return base unmodified
    return base;
  }
 
  // Evaluate constant
  if (base->type() == CONSTANT && power->type() == CONSTANT) {
    return constant_ptr(pow(base->val, power->val));
  }
 
  if (power->is_constant(Scalar(2))) {
    if (base->type() == LINEAR) {
      return make_expression_ptr<MultExpression<Scalar, QUADRATIC>>(base, base);
    } else {
      return make_expression_ptr<MultExpression<Scalar, NONLINEAR>>(base, base);
    }
  }
 
  return make_expression_ptr<PowExpression<Scalar, NONLINEAR>>(base, power);
}
 
template <typename Scalar>
struct SignExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr SignExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    if (x->val < Scalar(0)) {
      return Scalar(-1);
    } else if (x->val == Scalar(0)) {
      return Scalar(0);
    } else {
      return Scalar(1);
    }
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "sign"; }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sign(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    if (x->val < Scalar(0)) {
      return constant_ptr(Scalar(-1));
    } else if (x->val == Scalar(0)) {
      // Return zero
      return x;
    } else {
      return constant_ptr(Scalar(1));
    }
  }
 
  return make_expression_ptr<SignExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SinExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr SinExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::sin;
    return sin(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "sin"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::cos;
    return this->adj * cos(x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * cos(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sin(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::sin;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(sin(x->val));
  }
 
  return make_expression_ptr<SinExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SinhExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr SinhExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::sinh;
    return sinh(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "sinh"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::cosh;
    return this->adj * cosh(x->val);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const { return this->adj_expr * cosh(x); }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sinh(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::sinh;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(sinh(x->val));
  }
 
  return make_expression_ptr<SinhExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SqrtExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr SqrtExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::sqrt;
    return sqrt(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "sqrt"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::sqrt;
    return this->adj / (Scalar(2) * sqrt(x->val));
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    return this->adj_expr / (constant_ptr(Scalar(2)) * sqrt(x));
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sqrt(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::sqrt;
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    if (x->val == Scalar(0)) {
      // Return zero
      return x;
    } else if (x->val == Scalar(1)) {
      return x;
    } else {
      return constant_ptr(sqrt(x->val));
    }
  }
 
  return make_expression_ptr<SqrtExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct TanExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr TanExpression(ExpressionPtr<Scalar> x) : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::tan;
    return tan(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "tan"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::cos;
 
    auto c = cos(x->val);
    return this->adj / (c * c);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    auto c = cos(x);
    return this->adj_expr / (c * c);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> tan(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::tan;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(tan(x->val));
  }
 
  return make_expression_ptr<TanExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct TanhExpression final : Expression<Scalar> {
  ExpressionPtr<Scalar> x;
 
  explicit constexpr TanhExpression(ExpressionPtr<Scalar> x)
      : x{std::move(x)} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const override {
    func(x.get());
  }
 
  Scalar value() const override {
    using std::tanh;
    return tanh(x->val);
  }
 
  ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
  std::string_view name() const override { return "tanh"; }
 
  void accumulate_adjoints() const override { x->adj += grad_l(); }
 
  void accumulate_adjoints_expr() const override {
    x->adj_expr += grad_expr_l();
  }
 
 private:
  Scalar grad_l() const {
    using std::cosh;
 
    auto c = cosh(x->val);
    return this->adj / (c * c);
  }
 
  ExpressionPtr<Scalar> grad_expr_l() const {
    auto c = cosh(x);
    return this->adj_expr / (c * c);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> tanh(const ExpressionPtr<Scalar>& x) {
  using enum ExpressionType;
  using std::tanh;
 
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    // Return zero, which x currently is
    return x;
  }
 
  // Evaluate constant
  if (x->type() == CONSTANT) {
    return constant_ptr(tanh(x->val));
  }
 
  return make_expression_ptr<TanhExpression<Scalar>>(x);
}
 
}  // namespace slp::detail