Compiler support for laplace_latent_solve(_tol)

src/frontend/Semantic_error.ml

@@ -491,6 +491,7 @@ end
 
 module ExpressionError = struct
   type t =
+    | InvalidLaplaceLatentSolveFunction
     | InvalidSizeDeclRng
     | InvalidRngFunction
     | InvalidUnnormalizedFunction of string
@@ -510,6 +511,10 @@ module ExpressionError = struct
     | IllTypedPostfixOperator of Operator.t * UnsizedType.t
 
   let pp ppf = function
+    | InvalidLaplaceLatentSolveFunction ->
+        Fmt.text ppf
+          "Functions laplace_latent_solve and laplace_latent_solve_tol are \
+           only allowed in the generated quantities block."
     | InvalidSizeDeclRng ->
         Fmt.pf ppf
           "Random number generators are not allowed in top level size \
@@ -969,6 +974,9 @@ let ident_has_unnormalized_suffix loc name =
 let invalid_decl_rng_fn loc =
   (loc, ExpressionError ExpressionError.InvalidSizeDeclRng)
 
+let invalid_laplace_latent_solve_fn loc =
+  (loc, ExpressionError ExpressionError.InvalidLaplaceLatentSolveFunction)
+
 let invalid_rng_fn loc =
   (loc, ExpressionError ExpressionError.InvalidRngFunction)
 

src/frontend/Semantic_error.mli

@@ -164,6 +164,7 @@ val ident_not_in_scope :
   -> t
 
 val invalid_decl_rng_fn : Location_span.t -> t
+val invalid_laplace_latent_solve_fn : Location_span.t -> t
 val invalid_rng_fn : Location_span.t -> t
 val invalid_unnormalized_fn : Location_span.t -> string -> t
 val udf_is_unnormalized_fn : Location_span.t -> string -> t

src/frontend/Typechecker.ml

@@ -496,6 +496,16 @@ let verify_fn_rng cf loc id =
        || cf.current_block = TData)
     then Semantic_error.invalid_rng_fn loc |> error
 
+let is_laplace_latent_solve name =
+  Stan_math_signatures.is_embedded_laplace_fn name
+  && String.is_substring name ~substring:"_solve"
+
+(** Laplace latent solve can only be used in Generated Quantities block. *)
+let verify_laplace_latent_solve cf loc id =
+  if is_laplace_latent_solve id.name then
+    if cf.current_block <> GQuant then
+      Semantic_error.invalid_laplace_latent_solve_fn loc |> error
+
 let mk_fun_app ~is_cond_dist ~loc kind name args ~type_ : Ast.typed_expression =
   let fn =
     if is_cond_dist then CondDistApp (kind, name, args)
@@ -919,6 +929,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))
@@ -973,6 +985,7 @@ and check_funapp loc cf tenv ~is_cond_dist id (es : Ast.typed_expression list) =
   verify_fn_target_plus_equals cf loc id;
   verify_fn_jacobian_plus_equals cf loc tenv id es;
   verify_fn_rng cf loc id;
+  verify_laplace_latent_solve cf loc id;
   verify_unnormalized cf loc id;
   res
 

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_model.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;
+}
+
+model {
+   // laplace_latent_solve is only callable in Generated Quantities block
+   tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
+      = 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_model.stan
+Semantic error in 'bad_solve_model.stan', line 38, column 8 to line 39, column 64:
+   -------------------------------------------------
+    36:     // laplace_latent_solve is only callable in Generated Quantities block
+    37:     tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
+    38:        = laplace_latent_solve(ll_function, (eta, log_ye, y), 1,
+                 ^
+    39:                               K_function, (x, n_obs, alpha, rho));
+    40:  }
+   -------------------------------------------------
+
+Functions laplace_latent_solve and laplace_latent_solve_tol are only allowed
+in the generated quantities
+block.
+[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: