Unexpected error when passing a multi-index to a offset or multiply constraint

[WardBrian]

Model and data:

data {
  int<lower=1> n_group;
  int<lower=2> n_a;
  array[n_a] int group_index_for_a;
}
parameters {
  vector<lower=0>[n_group] a_group_mu;
  // can also try offset instead
  vector<multiplier=a_group_mu[group_index_for_a]>[n_a] a_effect;
}

{
  "n_group": 5,
  "n_a": 4,
  "group_index_for_a": [5, 3, 1, 5]
}

When using offset, this prints:

Unrecoverable error evaluating the log probability at the initial value.
Exception: constraint: a_effect (140728813511656, 1) and offset (4, 1) must match in size

(note that the size is essentially random and changes each run!!)

When using multiplier as above, it simply segfaults in Eigen::internal::binary_evaluator<Eigen::CwiseBinaryOp<Eigen::internal::scalar_product_op<double, double>, with the backtrace
suggesting the issue is in write_array (meaning the prim overload of everything)

cc @mike-lawrence

[WardBrian]

@SteveBronder It seems this was exposed by stan-dev/stanc3#1441 but I'm not sure. I'm also not sure if it would be considered a bug in the codegen or in the math library

Could you take a look?

[SteveBronder]

When I run this locally I am getting an error that the multiplier is infinite at some values.

Iteration: 2000 / 2000 [100%]  (Sampling)
Informational Message: The current Metropolis proposal is about to be rejected because of the following issue:
Exception: offset_multiplier_constrain: multiplier[2] is inf, but must be positive finite! (in 'examples/test.stan', line 9, column 2 to column 65)
If this warning occurs sporadically, such as for highly constrained variable types like covariance matrices, then the sampler is fine,
but if this warning occurs often then your model may be either severely ill-conditioned or misspecified.


 Elapsed Time: 0.144 seconds (Warm-up)
               0.175 seconds (Sampling)
               0.319 seconds (Total)

So looking at the code

      auto a_effect =
        in__.template read_constrain_offset_multiplier<
          Eigen::Matrix<local_scalar_t__,-1,1>, jacobian__>(0,
          stan::model::rvalue(a_group_mu, "a_group_mu",
            stan::model::index_multi(group_index_for_a)), lp__, n_a);

Here constrain_offset_multiplier is calling fma under the hood

template <typename T, typename M, typename S,
          require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
              T, M, S>* = nullptr>
inline auto offset_multiplier_constrain(const T& x, const M& mu,
                                        const S& sigma) {
  const auto& mu_ref = to_ref(mu);
  const auto& sigma_ref = to_ref(sigma);
  if (is_matrix<T>::value && is_matrix<M>::value) {
    check_matching_dims("offset_multiplier_constrain", "x", x, "mu", mu);
  }
  if (is_matrix<T>::value && is_matrix<S>::value) {
    check_matching_dims("offset_multiplier_constrain", "x", x, "sigma", sigma);
  } else if (is_matrix<M>::value && is_matrix<S>::value) {
    check_matching_dims("offset_multiplier_constrain", "mu", mu, "sigma",
                        sigma);
  }

  check_finite("offset_multiplier_constrain", "offset", value_of_rec(mu_ref));
  check_positive_finite("offset_multiplier_constrain", "multiplier",
                        value_of_rec(sigma_ref));
  return fma(sigma_ref, x, mu_ref);
}

template <typename T1, typename T2, typename T3,
          require_any_matrix_t<T1, T2, T3>* = nullptr,
          require_not_var_t<return_type_t<T1, T2, T3>>* = nullptr>
inline auto fma(T1&& x, T2&& y, T3&& z) {
  if (is_matrix<T1>::value && is_matrix<T2>::value) {
    check_matching_dims("fma", "x", x, "y", y);
  }
  if (is_matrix<T1>::value && is_matrix<T3>::value) {
    check_matching_dims("fma", "x", x, "z", z);
  } else if (is_matrix<T2>::value && is_matrix<T3>::value) {
    check_matching_dims("fma", "y", y, "z", z);
  }
  return make_holder(
      [](auto&& x, auto&& y, auto&& z) {
        return ((as_array_or_scalar(x) * as_array_or_scalar(y))
                + as_array_or_scalar(z))
            .matrix();
      },
      std::forward<T1>(x), std::forward<T2>(y), std::forward<T3>(z));
}

I think what might be happening is that we are making an rvalue eigen expression here when making the multi index.

stan::model::rvalue(a_group_mu, "a_group_mu",
            stan::model::index_multi(group_index_for_a))

Then the fma creates a Holder class storing that expression. I have to check, but I think that when the holder class is returned the multi index would be out of scope. I think a fix here is to make parameters that have constraints no longer use auto. At some point I can fix this by going into the math library and adding perfect forwarding everywhere. But we shouldn't wait for that when I think just removing auto for constrained params is a much easier fix.