Feature/issue 3249 minmax rev

stan/math/prim/fun/max.hpp

@@ -21,7 +21,8 @@ namespace math {
  * @param y second argument
  * @return maximum value of the two arguments
  */
-template <typename T1, typename T2, require_all_arithmetic_t<T1, T2>* = nullptr>
+template <typename T1, typename T2, require_all_arithmetic_t<T1, T2>* = nullptr,
+          require_all_not_var_t<T1, T2>* = nullptr>
 inline auto max(T1 x, T2 y) {
   return std::max(x, y);
 }
@@ -38,7 +39,8 @@ inline auto max(T1 x, T2 y) {
  * @throws <code>std::invalid_argument</code> if the vector is size zero and the
  * scalar type in the container is integer
  */
-template <typename T, require_container_t<T>* = nullptr>
+template <typename T, require_container_t<T>* = nullptr,
+          require_not_st_var<T>* = nullptr>
 inline value_type_t<T> max(T&& m) {
   if constexpr (std::is_integral<value_type_t<T>>::value) {
     check_nonzero_size("max", "int vector", m);

stan/math/prim/fun/min.hpp

@@ -37,7 +37,8 @@ inline auto min(T1 x, T2 y) {
  * @throws <code>std::invalid_argument</code> if the vector is size zero and the
  * scalar type in the container is integer
  */
-template <typename T, require_container_t<T>* = nullptr>
+template <typename T, require_container_t<T>* = nullptr,
+          require_not_st_var<T>* = nullptr>
 inline value_type_t<T> min(T&& m) {
   if constexpr (std::is_integral<value_type_t<T>>::value) {
     check_nonzero_size("min", "int vector", m);

stan/math/rev/fun.hpp

@@ -118,10 +118,12 @@
 #include <stan/math/rev/fun/logit.hpp>
 #include <stan/math/rev/fun/matrix_exp_multiply.hpp>
 #include <stan/math/rev/fun/matrix_power.hpp>
+#include <stan/math/rev/fun/max.hpp>
 #include <stan/math/rev/fun/mdivide_left.hpp>
 #include <stan/math/rev/fun/mdivide_left_ldlt.hpp>
 #include <stan/math/rev/fun/mdivide_left_spd.hpp>
 #include <stan/math/rev/fun/mdivide_left_tri.hpp>
+#include <stan/math/rev/fun/min.hpp>
 #include <stan/math/rev/fun/modified_bessel_first_kind.hpp>
 #include <stan/math/rev/fun/modified_bessel_second_kind.hpp>
 #include <stan/math/rev/fun/multiply.hpp>

stan/math/rev/fun/max.hpp

@@ -0,0 +1,129 @@
+#ifndef STAN_MATH_REV_FUN_MAX_HPP
+#define STAN_MATH_REV_FUN_MAX_HPP
+
+#include <stan/math/rev/core.hpp>
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/fun/is_nan.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/core.hpp>
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/is_nan.hpp>
+#include <stan/math/prim/fun/max.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Returns the maximum value of the two specified scalar arguments,
+ * where at least one argument is a 'var'.
+ *
+ * @tparam Scal1 type of first argument (must be 'var' if Scal2 is not)
+ * @tparam Scal2 type of second argument (must be 'var' if Scal1 is not)
+ * @param[in] x first argument
+ * @param[in] y second argument
+ * @return maximum value of the two arguments
+ */
+template <typename Scal1, typename Scal2,
+          require_any_var_t<Scal1, Scal2>* = nullptr>
+inline var max(Scal1 x, Scal2 y) {
+  double x_val = stan::math::value_of(x);
+  double y_val = stan::math::value_of(y);
+  if (unlikely(is_nan(x_val))) {
+    if (unlikely(is_nan(y_val))) {
+      return make_callback_var(NOT_A_NUMBER, [x, y](auto& vi) mutable {
+        if constexpr (is_var<Scal1>::value)
+          x.adj() = NOT_A_NUMBER;
+        if constexpr (is_var<Scal2>::value)
+          y.adj() = NOT_A_NUMBER;
+      });
+    }
+    return y;
+  }
+  if (unlikely(is_nan(y_val))) {
+    return x;
+  }
+  double max_val = std::fmax(x_val, y_val);
+  return make_callback_var(max_val, [x, y, x_val, y_val](auto& vi) mutable {
+    if (x_val > y_val) {
+      if constexpr (is_var<Scal1>::value)
+        x.adj() += vi.adj();
+    } else if (y_val > x_val) {
+      if constexpr (is_var<Scal2>::value)
+        y.adj() += vi.adj();
+    } else {
+      if constexpr (is_var<Scal1>::value)
+        x.adj() += vi.adj() * 0.5;
+      if constexpr (is_var<Scal2>::value)
+        y.adj() += vi.adj() * 0.5;
+    }
+  });
+}
+/**
+ * Return the maximum value in a container of 'var'.
+ * @tparam T Type of container with 'var' scalar type
+ * @param[in] x container (Eigen matrix, vector, row vector or std vector)
+ * @return maximum value in the container, or -infinity if size is zero
+ */
+template <typename T, require_st_var<T>* = nullptr,
+          require_container_t<T>* = nullptr,
+          require_not_var_matrix_t<T>* = nullptr>
+inline var max(T&& x) {
+  if (unlikely(x.size() == 0)) {
+    return NEGATIVE_INFTY;
+  }
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& v) {
+    const auto& x_ref = to_ref(v);
+    arena_t<std::decay_t<decltype(x_ref)>> x_arena(x_ref);
+    double max_val = x_arena.val().maxCoeff();
+    return make_callback_var(max_val, [x_arena, max_val](auto& vi) mutable {
+      if (unlikely(is_nan(max_val))) {
+        for (Eigen::Index i = 0; i < x_arena.size(); ++i) {
+          if (is_nan(x_arena.val().coeff(i))) {
+            x_arena.adj().coeffRef(i) = NOT_A_NUMBER;
+          }
+        }
+        return;
+      }
+      auto is_max = (x_arena.val().array() == max_val);
+      double count = is_max.template cast<double>().sum();
+      double adj_to_propagate = vi.adj() / count;
+      x_arena.adj().array()
+          += is_max.template cast<double>() * adj_to_propagate;
+    });
+  });
+}
+
+/**
+ * Return the maximum value in a `var_value` containing an Eigen matrix,
+ * vector, or row vector.
+ *
+ * @tparam T Type of `var_value` with an Eigen inner type
+ * @param[in] x `var_value` matrix
+ * @return maximum value in the matrix, or -infinity if size is zero
+ */
+template <typename T, require_var_matrix_t<T>* = nullptr>
+inline var max(const T& x) {
+  if (unlikely(x.size() == 0)) {
+    return NEGATIVE_INFTY;
+  }
+  double max_val = x.val().maxCoeff();
+  return make_callback_var(max_val, [x, max_val](auto& vi) mutable {
+    if (unlikely(is_nan(max_val))) {
+      for (Eigen::Index i = 0; i < x.size(); ++i) {
+        if (is_nan(x.val().coeff(i))) {
+          x.adj().coeffRef(i) = NOT_A_NUMBER;
+        }
+      }
+      return;
+    }
+    auto is_max = (x.val().array() == max_val);
+    double count = is_max.template cast<double>().sum();
+    double adj_to_propagate = vi.adj() / count;
+    x.adj().array() += is_max.template cast<double>() * adj_to_propagate;
+  });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/fun/min.hpp

@@ -0,0 +1,131 @@
+#ifndef STAN_MATH_REV_FUN_MIN_HPP
+#define STAN_MATH_REV_FUN_MIN_HPP
+
+#include <stan/math/rev/core.hpp>
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/fun/is_nan.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/core.hpp>
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/is_nan.hpp>
+#include <stan/math/prim/fun/min.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Returns the minimum value of the two specified scalar arguments,
+ * where at least one argument is a 'var'.
+ *
+ * @tparam Scal1 type of first argument (must be 'var' if Scal2 is not)
+ * @tparam Scal2 type of second argument (must be 'var' if Scal1 is not)
+ * @param[in] x first argument
+ * @param[in] y second argument
+ * @return minimum value of the two arguments
+ */
+template <typename Scal1, typename Scal2,
+          require_any_var_t<Scal1, Scal2>* = nullptr>
+inline var min(Scal1 x, Scal2 y) {
+  double x_val = stan::math::value_of(x);
+  double y_val = stan::math::value_of(y);
+  if (unlikely(is_nan(x_val))) {
+    if (unlikely(is_nan(y_val))) {
+      return make_callback_var(NOT_A_NUMBER, [x, y](auto& vi) mutable {
+        if constexpr (is_var<Scal1>::value)
+          x.adj() = NOT_A_NUMBER;
+        if constexpr (is_var<Scal2>::value)
+          y.adj() = NOT_A_NUMBER;
+      });
+    }
+    return y;
+  }
+  if (unlikely(is_nan(y_val))) {
+    return x;
+  }
+  double min_val = std::fmin(x_val, y_val);
+  return make_callback_var(min_val, [x, y, x_val, y_val](auto& vi) mutable {
+    if (x_val < y_val) {
+      if constexpr (is_var<Scal1>::value)
+        x.adj() += vi.adj();
+    } else if (y_val < x_val) {
+      if constexpr (is_var<Scal2>::value)
+        y.adj() += vi.adj();
+    } else {
+      if constexpr (is_var<Scal1>::value)
+        x.adj() += vi.adj() * 0.5;
+      if constexpr (is_var<Scal2>::value)
+        y.adj() += vi.adj() * 0.5;
+    }
+  });
+}
+
+/**
+ * Return the minimum value in a container of 'var'.
+ *
+ * @tparam T Type of container with 'var' scalar type
+ * @param[in] x container (Eigen matrix, vector, row vector or std vector)
+ * @return minimum value in the container, or positive infinity if size is zero
+ */
+template <typename T, require_st_var<T>* = nullptr,
+          require_container_t<T>* = nullptr,
+          require_not_var_matrix_t<T>* = nullptr>
+inline var min(T&& x) {
+  if (unlikely(x.size() == 0)) {
+    return INFTY;
+  }
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& v) {
+    const auto& x_ref = to_ref(v);
+    arena_t<std::decay_t<decltype(x_ref)>> x_arena(x_ref);
+    double min_val = x_arena.val().minCoeff();
+    return make_callback_var(min_val, [x_arena, min_val](auto& vi) mutable {
+      if (unlikely(is_nan(min_val))) {
+        for (Eigen::Index i = 0; i < x_arena.size(); ++i) {
+          if (is_nan(x_arena.val().coeff(i))) {
+            x_arena.adj().coeffRef(i) = NOT_A_NUMBER;
+          }
+        }
+        return;
+      }
+      auto is_min = (x_arena.val().array() == min_val);
+      double count = is_min.template cast<double>().sum();
+      double adj_to_propagate = vi.adj() / count;
+      x_arena.adj().array()
+          += is_min.template cast<double>() * adj_to_propagate;
+    });
+  });
+}
+
+/**
+ * Return the minimum value in a `var_value` containing an Eigen matrix,
+ * vector, or row vector.
+ *
+ * @tparam T Type of `var_value` with an Eigen inner type
+ * @param[in] x `var_value` matrix
+ * @return minimum value in the matrix, or infinity if size is zero
+ */
+template <typename T, require_var_matrix_t<T>* = nullptr>
+inline var min(const T& x) {
+  if (unlikely(x.size() == 0)) {
+    return INFTY;
+  }
+  double min_val = x.val().minCoeff();
+  return make_callback_var(min_val, [x, min_val](auto& vi) mutable {
+    if (unlikely(is_nan(min_val))) {
+      for (Eigen::Index i = 0; i < x.size(); ++i) {
+        if (is_nan(x.val().coeff(i))) {
+          x.adj().coeffRef(i) = NOT_A_NUMBER;
+        }
+      }
+      return;
+    }
+    auto is_min = (x.val().array() == min_val);
+    double count = is_min.template cast<double>().sum();
+    double adj_to_propagate = vi.adj() / count;
+    x.adj().array() += is_min.template cast<double>() * adj_to_propagate;
+  });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

test/unit/math/mix/fun/max_test.cpp

@@ -13,25 +13,48 @@ inline void expect_max(const T& m) {
   stan::test::expect_ad(f, v);
   stan::test::expect_ad(f, rv);
   stan::test::expect_ad(f, m);
+  stan::test::expect_ad_matvar(f, v);
+  stan::test::expect_ad_matvar(f, rv);
+  stan::test::expect_ad_matvar(f, m);
 }
 
 TEST(MathMixMatFun, max) {
   Eigen::MatrixXd a(0, 0);
   expect_max(a);
-
   Eigen::MatrixXd b(1, 1);
   b << 12;
   expect_max(b);
-
   Eigen::MatrixXd c(3, 1);
   c << 100, 0, -3;
   expect_max(c);
-
   Eigen::MatrixXd d(1, 3);
   d << 100, 0, -3;
   expect_max(d);
-
   Eigen::MatrixXd e(3, 2);
   e << -100, 0, 1, 20, -40, 2;
   expect_max(e);
+  Eigen::MatrixXd ties(2, 2);
+  ties << 10.5, 1.0, 10.5, 10.5;
+  expect_max(ties);
+  double inf = std::numeric_limits<double>::infinity();
+  Eigen::VectorXd o(3);
+  o << -inf, 5.0, inf;
+  expect_max(o);
+  double nan = std::numeric_limits<double>::quiet_NaN();
+  Eigen::VectorXd n1(3);
+  n1 << 1.0, nan, 2.0;
+  expect_max(n1);
+  Eigen::MatrixXd n2(2, 2);
+  n2 << nan, nan, nan, nan;
+  expect_max(n2);
+}
+TEST(MathMixMatFun, max_binary) {
+  auto f = [](const auto& x, const auto& y) { return stan::math::max(x, y); };
+  stan::test::expect_ad(f, 1.0, 2.0);
+  stan::test::expect_ad(f, 2.0, 1.0);
+  stan::test::expect_ad(f, 3.0, 3.0);
+  double nan = std::numeric_limits<double>::quiet_NaN();
+  stan::test::expect_ad(f, nan, 1.0);
+  stan::test::expect_ad(f, 1.0, nan);
+  stan::test::expect_ad(f, nan, nan);
 }