Compiler support for laplace_latent_solve(_tol)

[florence-bockting]

Submission Checklist

This PR refers to math PR#337 in which we implemented laplace_latent_solve and laplace_latent_solve_tol.

• Run unit tests
• Documentation
  • If a user-facing facing change was made, the documentation PR is here: Document laplace_latent_solve(_tol) docs#968
  • OR, no user-facing changes were made

Release notes

• Adds Stan compiler support for Math’s laplace_latent_solve(_tol), which return tuple(vector, matrix) (posterior mean and Cholesky factor).

Copyright and Licensing

By submitting this pull request, the copyright holder is agreeing to
license the submitted work under the BSD 3-clause license (https://opensource.org/licenses/BSD-3-Clause)

This PR used AI assistance (Cursor).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.51%. Comparing base (5496e41) to head (ec9a2f8).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1649   +/-   ##
=======================================
  Coverage   91.51%   91.51%           
=======================================
  Files          67       67           
  Lines        9977     9979    +2     
=======================================
+ Hits         9130     9132    +2     
  Misses        847      847           

Files with missing lines	Coverage Δ
src/frontend/Typechecker.ml	94.88% <100.00%> (+<0.01%)	⬆️
src/stan_math_signatures/Stan_math_signatures.ml	86.27% <ø> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[WardBrian]

Thanks @florence-bockting -- question about 6e2c28b -- is it actually a requirement that these functions only be called in generated quantities? If so we will want to enforce that in the typechecker

[florence-bockting]

Thanks @florence-bockting -- question about 6e2c28b -- is it actually a requirement that these functions only be called in generated quantities? If so we will want to enforce that in the typechecker

Thank you for pointing me towards the typechecker!
Short about the background: As laplace_latent_solve currently only goes through laplace_base_rng it inherits the constraint on CovarArgs and LLArgs:

template <bool ReturnMeanAndCovCholesky = false, typename LLFunc,
          typename LLArgs, typename CovarFun, typename CovarArgs,
          bool InitTheta, typename RNG,
          require_t<is_all_arithmetic_scalar<CovarArgs, LLArgs>>* = nullptr>

For the model block this fails as some of the arguments are variables. Here an example snippet of a corresponding model and generated quantities block:

model {
  rho ~ inv_gamma(rho_location_prior, rho_scale_prior);
  alpha ~ inv_gamma(alpha_location_prior, alpha_scale_prior);
  eta ~ normal(0, 1);

  target += laplace_marginal(ll_function, (eta, log_ye, y), hessian_block_size,
                             K_function, (x, n_obs, alpha, rho));

  target += laplace_marginal_tol(ll_function, (eta, log_ye, y), hessian_block_size,
                                 K_function, (x, n_obs, alpha, rho),
                                (theta_0, tolerance, max_num_steps,
                                solver, max_steps_line_search, allow_fallthrough));
}
generated quantities {
  tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
      = laplace_latent_solve(ll_function, (eta, log_ye, y), hessian_block_size,
                             K_function, (x, n_obs, alpha, rho));
}

I am not sure whether there are use cases where a user wants to call laplace_latent_solve in the model block (@avehtari )? In this case, we would have to change the definition in math. Otherwise I would add the corresponding constraint in the Typechecker.ml.

[SteveBronder]

So since we never actually use the rng for laplace_latent_solve idt there is any reason we could not have it in the model block. So language wise I don't see an issue arising because of a user calling this in the model block.

[florence-bockting]

So since we never actually use the rng for laplace_latent_solve idt there is any reason we could not have it in the model block. So language wise I don't see an issue arising because of a user calling this in the model block.

The definition of laplace_latent_solve is as follows:

template <typename LLFunc, typename LLArgs, typename CovarFun,
          typename CovarArgs>
inline auto laplace_latent_solve(LLFunc&& ll_fun, LLArgs&& ll_args,
                                 int hessian_block_size,
                                 CovarFun&& covariance_function,
                                 CovarArgs&& covar_args, std::ostream* msgs) {
  auto options = laplace_options_default{hessian_block_size};
  internal::laplace_unused_rng unused_rng;
  return laplace_base_rng<true>(std::forward<LLFunc>(ll_fun),
                                std::forward<LLArgs>(ll_args),
                                std::forward<CovarFun>(covariance_function),
                                std::forward<CovarArgs>(covar_args),
                                std::move(options), unused_rng, msgs);
}

So, we actually call laplace_base_rng and therefore inherit the constraints on the arguments.

An alternative, to circumvent calling laplace_base_rng, could be to factor out the code for computing the mean and variance in an extra function and then to call this function in laplace_base_rng and in laplace_latent_solve.

On the other hand, I talked yesterday with Aki and he said, he could only come up with theoretical examples for calling laplace_latent_solve in the model block and thus would be fine to restrict it to the generated quantities block.

I don't know whether you prefer to keep it simple and leave it as is (with the missing Typechecker) or redefining laplace_latent_solve such that it generalizes to more situations.