Fix and test autodiff containers

Author: andrjohns

stan/math/fwd/fun/hypergeometric_pFq.hpp

@@ -7,6 +7,8 @@
 #include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/grad_pFq.hpp>
 #include <stan/math/prim/fun/hypergeometric_pFq.hpp>
+#include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 
 namespace stan {
 namespace math {
@@ -30,33 +32,27 @@ template <typename Ta, typename Tb, typename Tz,
           bool grad_z = !is_constant<Tz>::value,
           require_all_vector_t<Ta, Tb>* = nullptr,
           require_fvar_t<FvarT>* = nullptr>
-inline FvarT hypergeometric_pFq(const Ta& a, const Tb& b, const Tz& z) {
-  using PartialsT = partials_type_t<FvarT>;
-  using ARefT = ref_type_t<Ta>;
-  using BRefT = ref_type_t<Tb>;
-
-  ARefT a_ref = a;
-  BRefT b_ref = b;
+inline FvarT hypergeometric_pFq(Ta&& a, Tb&& b, Tz&& z) {
+  auto&& a_ref = to_ref(as_column_vector_or_scalar(a));
+  auto&& b_ref = to_ref(as_column_vector_or_scalar(b));
   auto&& a_val = value_of(a_ref);
   auto&& b_val = value_of(b_ref);
   auto&& z_val = value_of(z);
-  PartialsT pfq_val = hypergeometric_pFq(a_val, b_val, z_val);
+
+  partials_type_t<FvarT> pfq_val = hypergeometric_pFq(a_val, b_val, z_val);
   auto grad_tuple
       = grad_pFq<grad_a, grad_b, grad_z>(pfq_val, a_val, b_val, z_val);
 
   FvarT rtn = FvarT(pfq_val, 0.0);
 
-  if (grad_a) {
-    rtn.d_ += dot_product(forward_as<promote_scalar_t<FvarT, ARefT>>(a_ref).d(),
-                          std::get<0>(grad_tuple));
+  if constexpr (grad_a) {
+    rtn.d_ += dot_product(a_ref.d(), std::get<0>(grad_tuple));
   }
-  if (grad_b) {
-    rtn.d_ += dot_product(forward_as<promote_scalar_t<FvarT, BRefT>>(b_ref).d(),
-                          std::get<1>(grad_tuple));
+  if constexpr (grad_b) {
+    rtn.d_ += dot_product(b_ref.d(), std::get<1>(grad_tuple));
   }
-  if (grad_z) {
-    rtn.d_ += forward_as<promote_scalar_t<FvarT, Tz>>(z).d_
-              * std::get<2>(grad_tuple);
+  if constexpr (grad_z) {
+    rtn.d_ += z.d_ * std::get<2>(grad_tuple);
   }
 
   return rtn;

stan/math/rev/fun/hypergeometric_pFq.hpp

@@ -3,6 +3,7 @@
 
 #include <stan/math/rev/core.hpp>
 #include <stan/math/rev/meta.hpp>
+#include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/grad_pFq.hpp>
 #include <stan/math/prim/fun/hypergeometric_pFq.hpp>
 
@@ -25,27 +26,24 @@ template <typename Ta, typename Tb, typename Tz,
           bool grad_a = !is_constant<Ta>::value,
           bool grad_b = !is_constant<Tb>::value,
           bool grad_z = !is_constant<Tz>::value,
-          require_all_matrix_t<Ta, Tb>* = nullptr,
+          require_all_vector_t<Ta, Tb>* = nullptr,
           require_return_type_t<is_var, Ta, Tb, Tz>* = nullptr>
-inline var hypergeometric_pFq(const Ta& a, const Tb& b, const Tz& z) {
-  arena_t<Ta> arena_a = a;
-  arena_t<Tb> arena_b = b;
+inline var hypergeometric_pFq(Ta&& a, Tb&& b, Tz&& z) {
+  auto&& arena_a = to_arena(as_column_vector_or_scalar(std::forward<Ta>(a)));
+  auto&& arena_b = to_arena(as_column_vector_or_scalar(std::forward<Tb>(b)));
   auto pfq_val = hypergeometric_pFq(arena_a.val(), arena_b.val(), value_of(z));
   return make_callback_var(
       pfq_val, [arena_a, arena_b, z, pfq_val](auto& vi) mutable {
         auto grad_tuple = grad_pFq<grad_a, grad_b, grad_z>(
             pfq_val, arena_a.val(), arena_b.val(), value_of(z));
         if constexpr (grad_a) {
-          forward_as<promote_scalar_t<var, Ta>>(arena_a).adj()
-              += vi.adj() * std::get<0>(grad_tuple);
+          arena_a.adj() += vi.adj() * std::get<0>(grad_tuple);
         }
         if constexpr (grad_b) {
-          forward_as<promote_scalar_t<var, Tb>>(arena_b).adj()
-              += vi.adj() * std::get<1>(grad_tuple);
+          arena_b.adj() += vi.adj() * std::get<1>(grad_tuple);
         }
         if constexpr (grad_z) {
-          forward_as<promote_scalar_t<var, Tz>>(z).adj()
-              += vi.adj() * std::get<2>(grad_tuple);
+          z.adj() += vi.adj() * std::get<2>(grad_tuple);
         }
       });
 }

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

@@ -2,6 +2,9 @@
 #include <limits>
 
 TEST(mathMixScalFun, hyper_2f2) {
+  using stan::math::to_array_1d;
+  using stan::math::to_row_vector;
+
   auto f = [](const auto& a, const auto& b, const auto& z) {
     using stan::math::hypergeometric_pFq;
     return hypergeometric_pFq(a, b, z);
@@ -14,6 +17,9 @@ TEST(mathMixScalFun, hyper_2f2) {
   double z = 4;
 
   stan::test::expect_ad(f, in1, in2, z);
+  stan::test::expect_ad(f, to_array_1d(in1), to_row_vector(in2), z);
+  stan::test::expect_ad(f, to_row_vector(in1), to_array_1d(in2), z);
+  stan::test::expect_ad(f, to_array_1d(in1), to_array_1d(in2), z);
 }
 
 TEST(mathMixScalFun, hyper_2f3) {