Compiler support for laplace_latent_solve(_tol)

src/frontend/Typechecker.ml

@@ -919,6 +919,8 @@ and check_laplace_fn ~is_cond_dist loc cf tenv id tes =
         lik_args @ (hbs_arg :: cov_fun_type :: cov_tupl :: control_args) in
       let return_type =
         if String.is_suffix id.name ~suffix:"_rng" then UnsizedType.UVector
+        else if String.is_substring id.name ~substring:"_solve" then
+          UnsizedType.UTuple [UVector; UMatrix]
         else UnsizedType.UReal in
       mk_fun_app ~is_cond_dist ~loc
         (StanLib (Fun_kind.suffix_from_name id.name))

src/stan_math_signatures/Stan_math_signatures.ml

@@ -153,7 +153,8 @@ let is_reduce_sum_fn f =
 
 let embedded_laplace_functions =
   [ (* general fns *) "laplace_marginal"; "laplace_marginal_tol"
-  ; "laplace_latent_rng"; "laplace_latent_tol_rng"; (* "helpers" *)
+  ; "laplace_latent_rng"; "laplace_latent_tol_rng"; "laplace_latent_solve"
+  ; "laplace_latent_solve_tol"; (* "helpers" *)
     "laplace_marginal_bernoulli_logit_lpmf"
   ; "laplace_marginal_tol_bernoulli_logit_lpmf"
   ; "laplace_marginal_neg_binomial_2_log_lpmf"

test/integration/bad/embedded_laplace/bad_solve_return.stan

@@ -0,0 +1,39 @@
+functions {
+  // specify negative binomial likelihood with mean offset
+  real ll_function(vector theta, // latent Gaussian
+                   real eta, vector log_ye, // mean offset
+                   array[] int y) {
+    // observed count
+    return neg_binomial_2_lpmf(y | exp(log_ye + theta), eta);
+  }
+
+  // specify covariance function
+  matrix K_function(array[] vector x, int n_obs, real alpha, real rho) {
+    matrix[n_obs, n_obs] K = gp_exp_quad_cov(x, alpha, rho);
+    for (i in 1 : n_obs)
+      K[i, i] += 1e-8;
+    return K;
+  }
+}
+data {
+  int n_obs;
+  int n_coordinates;
+  array[n_obs] int y;
+  vector[n_obs] ye;
+  array[n_obs] vector[n_coordinates] x;
+}
+
+transformed data {
+  vector[n_obs] log_ye = log(ye);
+}
+parameters {
+  real<lower=0> alpha;
+  real<lower=0> rho;
+  real<lower=0> eta;
+}
+
+generated quantities {
+  // laplace_latent_solve returns tuple(vector, matrix), not vector
+  vector[n_obs] theta = laplace_latent_solve(ll_function, (eta, log_ye, y), 1,
+                                             K_function, (x, n_obs, alpha, rho));
+}

test/integration/bad/embedded_laplace/bad_solve_tol.stan

@@ -0,0 +1,40 @@
+functions {
+  // specify negative binomial likelihood with mean offset
+  real ll_function(vector theta, // latent Gaussian
+                   real eta, vector log_ye, // mean offset
+                   array[] int y) {
+    // observed count
+    return neg_binomial_2_lpmf(y | exp(log_ye + theta), eta);
+  }
+
+  // specify covariance function
+  matrix K_function(array[] vector x, int n_obs, real alpha, real rho) {
+    matrix[n_obs, n_obs] K = gp_exp_quad_cov(x, alpha, rho);
+    for (i in 1 : n_obs)
+      K[i, i] += 1e-8;
+    return K;
+  }
+}
+data {
+  int n_obs;
+  int n_coordinates;
+  array[n_obs] int y;
+  vector[n_obs] ye;
+  array[n_obs] vector[n_coordinates] x;
+}
+
+transformed data {
+  vector[n_obs] log_ye = log(ye);
+}
+parameters {
+  real<lower=0> alpha;
+  real<lower=0> rho;
+  real<lower=0> eta;
+}
+
+generated quantities {
+  // _tol variant requires a trailing control-parameter tuple
+  tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
+      = laplace_latent_solve_tol(ll_function, (eta, log_ye, y), 1, K_function,
+                                 (x, n_obs, alpha, rho));
+}

test/integration/bad/embedded_laplace/stanc.expected

@@ -325,6 +325,35 @@ Semantic error in 'bad_overload.stan', line 2, column 7 to column 51:
    -------------------------------------------------
 
 Identifier "laplace_marginal_tol_neg_binomial_2_log_lpmf" clashes with a non-overloadable Stan Math library function.
+[exit 1]
+  $ stanc bad_solve_return.stan
+Semantic error in 'bad_solve_return.stan', line 37, column 2 to line 38, column 81:
+   -------------------------------------------------
+    35:  generated quantities {
+    36:    // laplace_latent_solve returns tuple(vector, matrix), not vector
+    37:    vector[n_obs] theta = laplace_latent_solve(ll_function, (eta, log_ye, y), 1,
+           ^
+    38:                                               K_function, (x, n_obs, alpha, rho));
+    39:  }
+   -------------------------------------------------
+
+Ill-typed assignment statement.
+Expected the right hand side to have a type matching the destination (vector).
+Instead found type tuple(vector, matrix).
+[exit 1]
+  $ stanc bad_solve_tol.stan
+Semantic error in 'bad_solve_tol.stan', line 38, column 8 to line 39, column 56:
+   -------------------------------------------------
+    36:    // _tol variant requires a trailing control-parameter tuple
+    37:    tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
+    38:        = laplace_latent_solve_tol(ll_function, (eta, log_ye, y), 1, K_function,
+                 ^
+    39:                                   (x, n_obs, alpha, rho));
+    40:  }
+   -------------------------------------------------
+
+Missing control parameter tuple at the end of the call to "laplace_latent_solve_tol".
+Expected a tuple of 6 arguments for the control parameters.
 [exit 1]
   $ stanc bad_theta0.stan
 Semantic error in 'bad_theta0.stan', line 43, column 9 to column 16: