Fix issue 3146 for multiply_log

stan/math/fwd/fun/lmultiply.hpp

@@ -4,27 +4,8 @@
 #include <stan/math/fwd/meta.hpp>
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/fun/log.hpp>
+#include <stan/math/fwd/fun/multiply_log.hpp>
 #include <stan/math/prim/fun/lmultiply.hpp>
 #include <cmath>
 
-namespace stan {
-namespace math {
-
-template <typename T>
-inline fvar<T> lmultiply(const fvar<T>& x1, const fvar<T>& x2) {
-  return fvar<T>(lmultiply(x1.val_, x2.val_),
-                 x1.d_ * log(x2.val_) + x1.val_ * x2.d_ / x2.val_);
-}
-
-template <typename T>
-inline fvar<T> lmultiply(double x1, const fvar<T>& x2) {
-  return fvar<T>(lmultiply(x1, x2.val_), x1 * x2.d_ / x2.val_);
-}
-
-template <typename T>
-inline fvar<T> lmultiply(const fvar<T>& x1, double x2) {
-  return fvar<T>(lmultiply(x1.val_, x2), x1.d_ * log(x2));
-}
-}  // namespace math
-}  // namespace stan
 #endif

stan/math/fwd/fun/multiply_log.hpp

@@ -4,6 +4,7 @@
 #include <stan/math/fwd/meta.hpp>
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/fun/log.hpp>
+#include <stan/math/fwd/fun/value_of_rec.hpp>
 #include <stan/math/prim/fun/multiply_log.hpp>
 #include <cmath>
 
@@ -12,19 +13,29 @@ namespace math {
 
 template <typename T>
 inline fvar<T> multiply_log(const fvar<T>& x1, const fvar<T>& x2) {
+  if (value_of_rec(x1) == 0.0 && value_of_rec(x2) == 0.0) {
+    return fvar<T>(0.0);
+  }
   return fvar<T>(multiply_log(x1.val_, x2.val_),
                  x1.d_ * log(x2.val_) + x1.val_ * x2.d_ / x2.val_);
 }
 
 template <typename T>
 inline fvar<T> multiply_log(double x1, const fvar<T>& x2) {
+  if (x1 == 0.0 && value_of_rec(x2) == 0.0) {
+    return fvar<T>(0.0);
+  }
   return fvar<T>(multiply_log(x1, x2.val_), x1 * x2.d_ / x2.val_);
 }
 
 template <typename T>
 inline fvar<T> multiply_log(const fvar<T>& x1, double x2) {
+  if (value_of_rec(x1) == 0.0 && x2 == 0.0) {
+    return fvar<T>(0.0);
+  }
   return fvar<T>(multiply_log(x1.val_, x2), x1.d_ * log(x2));
 }
+
 }  // namespace math
 }  // namespace stan
 #endif

stan/math/opencl/kernel_generator/elt_function_cl.hpp

@@ -365,8 +365,8 @@ ADD_BINARY_FUNCTION_WITH_INCLUDES(log_diff_exp,
                                   opencl_kernels::log_diff_exp_device_function)
 ADD_BINARY_FUNCTION_WITH_INCLUDES(
     multiply_log, stan::math::opencl_kernels::multiply_log_device_function)
-ADD_BINARY_FUNCTION_WITH_INCLUDES(
-    lmultiply, stan::math::opencl_kernels::lmultiply_device_function)
+// ADD_BINARY_FUNCTION_WITH_INCLUDES(
+//    lmultiply, stan::math::opencl_kernels::lmultiply_device_function)
 
 #undef ADD_BINARY_FUNCTION_WITH_INCLUDES
 #undef ADD_UNARY_FUNCTION_WITH_INCLUDES

stan/math/opencl/rev/lmultiply.hpp

@@ -1,46 +0,0 @@
-#ifndef STAN_MATH_OPENCL_REV_LMULTIPLY_HPP
-#define STAN_MATH_OPENCL_REV_LMULTIPLY_HPP
-#ifdef STAN_OPENCL
-
-#include <stan/math/opencl/rev/adjoint_results.hpp>
-#include <stan/math/opencl/kernel_generator.hpp>
-#include <stan/math/rev/core.hpp>
-#include <stan/math/rev/fun/adjoint_of.hpp>
-#include <stan/math/rev/fun/value_of.hpp>
-
-namespace stan {
-namespace math {
-
-/**
- * Returns the elementwise `lmultiply()` of the input.
- *
- * @tparam T_a type of first expression
- * @tparam T_b type of second expression
- * @param a first expression
- * @param b second expression
- *
- * @return Elementwise `lmultiply()` of the input.
- */
-template <typename T_a, typename T_b,
-          require_all_prim_or_rev_kernel_expression_t<T_a, T_b>* = nullptr,
-          require_any_var_t<T_a, T_b>* = nullptr,
-          require_any_not_stan_scalar_t<T_a, T_b>* = nullptr>
-inline var_value<matrix_cl<double>> lmultiply(T_a&& a, T_b&& b) {
-  arena_t<T_a> a_arena = std::forward<T_a>(a);
-  arena_t<T_b> b_arena = std::forward<T_b>(b);
-
-  return make_callback_var(
-      lmultiply(value_of(a_arena), value_of(b_arena)),
-      [a_arena, b_arena](const vari_value<matrix_cl<double>>& res) mutable {
-        adjoint_results(a_arena, b_arena) += expressions(
-            elt_multiply(res.adj(), log(value_of(b_arena))),
-            elt_multiply(res.adj(),
-                         elt_divide(value_of(a_arena), value_of(b_arena))));
-      });
-}
-
-}  // namespace math
-}  // namespace stan
-
-#endif
-#endif

stan/math/opencl/rev/multiply_log.hpp

@@ -10,6 +10,16 @@
 
 namespace stan {
 namespace math {
+namespace internalcl {
+template <bool Cond, typename T>
+inline decltype(auto) conditional_sum(T&& x) {
+  if constexpr (Cond) {
+    return sum(std::forward<T>(x));
+  } else {
+    return std::forward<T>(x);
+  }
+}
+}  // namespace internalcl
 
 /**
  * Returns the elementwise `multiply_log()` of the input.
@@ -32,10 +42,26 @@ inline var_value<matrix_cl<double>> multiply_log(T_a&& a, T_b&& b) {
   return make_callback_var(
       multiply_log(value_of(a_arena), value_of(b_arena)),
       [a_arena, b_arena](const vari_value<matrix_cl<double>>& res) mutable {
-        adjoint_results(a_arena, b_arena) += expressions(
-            elt_multiply(res.adj(), log(value_of(b_arena))),
-            elt_multiply(res.adj(),
-                         elt_divide(value_of(a_arena), value_of(b_arena))));
+        constexpr bool is_scalar_a = !is_matrix_v<T_a>;
+        constexpr bool is_scalar_b = !is_matrix_v<T_b>;
+        using internalcl::conditional_sum;
+        auto is_zero = value_of(a_arena) == 0.0 && value_of(b_arena) == 0.0;
+        if constexpr (is_var<T_a>::value && is_var<T_b>::value) {
+          a_arena.adj() += conditional_sum<is_scalar_a>(select(
+              is_zero, 0.0, elt_multiply(res.adj(), log(value_of(b_arena)))));
+          b_arena.adj() += conditional_sum<is_scalar_b>(select(
+              is_zero, 0.0,
+              elt_multiply(res.adj(),
+                           elt_divide(value_of(a_arena), value_of(b_arena)))));
+        } else if constexpr (is_var<T_a>::value) {
+          a_arena.adj() += conditional_sum<is_scalar_a>(select(
+              is_zero, 0.0, elt_multiply(res.adj(), log(value_of(b_arena)))));
+        } else if constexpr (is_var<T_b>::value) {
+          b_arena.adj() += conditional_sum<is_scalar_b>(select(
+              is_zero, 0.0,
+              elt_multiply(res.adj(),
+                           elt_divide(value_of(a_arena), value_of(b_arena)))));
+        }
       });
 }