Compiler support for laplace_latent_solve(_tol)

src/frontend/Semantic_error.ml

@@ -69,6 +69,7 @@ module TypeError = struct
     | IlltypedLaplaceHessianBlockSize of
         string * (UnsizedType.autodifftype * UnsizedType.t) option
     | IlltypedLaplaceTolArgs of string * SignatureMismatch.function_mismatch
+    | IlltypedLaplaceLatentSolveArgs of string
     | AmbiguousFunctionPromotion of
         string
         * UnsizedType.t list option
@@ -328,6 +329,11 @@ module TypeError = struct
         Fmt.pf ppf "@[<hov>The %s to %a@ must be@ %a.%a@]"
           (laplace_tolerance_arg_name n)
           quoted name expected_types [expected] found_type found
+    | IlltypedLaplaceLatentSolveArgs name ->
+        Fmt.pf ppf
+          "@[<hov>All arguments to %a must be a data-only@ if used outside the \
+           generated quantities@ block.@ %a@]"
+          quoted name SignatureMismatch.data_only_msg ()
     | AmbiguousFunctionPromotion (name, arg_tys, signatures) ->
         let pp_sig ppf (rt, args, _) =
           Fmt.pf ppf "@[<hov>(@[<hov>%a@]) => %a@]"
@@ -892,6 +898,9 @@ let illtyped_laplace_hessian_block_size_arg loc name arg_ty =
 let illtyped_laplace_tolerance_args loc name mismatch =
   (loc, TypeError (TypeError.IlltypedLaplaceTolArgs (name, mismatch)))
 
+let illtyped_laplace_latent_solve_args loc name =
+  (loc, TypeError (TypeError.IlltypedLaplaceLatentSolveArgs name))
+
 let ambiguous_function_promotion loc name arg_tys signatures =
   ( loc
   , TypeError (TypeError.AmbiguousFunctionPromotion (name, arg_tys, signatures))

src/frontend/Semantic_error.mli

@@ -125,6 +125,8 @@ val illtyped_laplace_hessian_block_size_arg :
 val illtyped_laplace_tolerance_args :
   Location_span.t -> string -> SignatureMismatch.function_mismatch -> t
 
+val illtyped_laplace_latent_solve_args : Location_span.t -> string -> t
+
 val nonreturning_fn_expected_returning_found :
   Location_span.t -> string -> Location_span.t option -> t
 

src/frontend/Typechecker.ml

@@ -497,6 +497,10 @@ 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"
+
 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)
@@ -877,6 +881,15 @@ and check_laplace_fn ~is_cond_dist loc cf tenv id tes =
     else
       (* likelihood callback check *)
       match tes with
+      | _ :: {emeta= {ad_level; loc; _}; expr} :: _
+        when is_laplace_latent_solve id.name
+             && UnsizedType.is_autodifftype ad_level ->
+          let es = match expr with TupleExpr es -> es | _ -> [] in
+          let loc =
+            List.find_map es ~f:(fun {emeta= {loc; ad_level; _}; _} ->
+                Option.some_if (UnsizedType.is_autodifftype ad_level) loc)
+            |> Option.value ~default:loc in
+          Semantic_error.illtyped_laplace_latent_solve_args loc id.name |> error
       | {expr= Variable lik_fun; _} :: lik_tupl :: tes ->
           let lik_fun, lik_tupl =
             (* adds the function name to the global list that is checked
@@ -907,6 +920,15 @@ and check_laplace_fn ~is_cond_dist loc cf tenv id tes =
         |> error in
   (* Check the remaining arguments: initial guess, covariance, and tolerances *)
   match rest with
+  | _ :: {emeta= {ad_level; loc; _}; expr} :: _
+    when is_laplace_latent_solve id.name && UnsizedType.is_autodifftype ad_level
+    ->
+      let es = match expr with TupleExpr es -> es | _ -> [] in
+      let loc =
+        List.find_map es ~f:(fun {emeta= {loc; ad_level; _}; _} ->
+            Option.some_if (UnsizedType.is_autodifftype ad_level) loc)
+        |> Option.value ~default:loc in
+      Semantic_error.illtyped_laplace_latent_solve_args loc id.name |> error
   | {expr= Variable cov_fun; _} :: cov_tupl :: control_args ->
       let cov_fun_type, cov_tupl =
         check_function_callable_with_tuple cf tenv id cov_fun cov_tupl
@@ -920,6 +942,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_cov.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, (1.0, log_ye, y), 1,
+                             K_function, (x, n_obs, alpha, rho));
+}

test/integration/bad/embedded_laplace/bad_solve_lik.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,51 @@ 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_cov.stan
+Semantic error in 'bad_solve_cov.stan', line 39, column 52 to column 57:
+   -------------------------------------------------
+    37:     tuple(vector[n_obs], matrix[n_obs, n_obs]) mean_chol
+    38:        = laplace_latent_solve(ll_function, (1.0, log_ye, y), 1,
+    39:                               K_function, (x, n_obs, alpha, rho));
+                                                             ^
+    40:  }
+   -------------------------------------------------
+
+All arguments to "laplace_latent_solve" must be a data-only
+if used outside the generated quantities block. (Local variables are assumed
+to depend on parameters; same goes for function inputs unless they are marked
+with the keyword "data".)
+[exit 1]
+  $ stanc bad_solve_lik.stan
+Semantic error in 'bad_solve_lik.stan', line 38, column 43 to column 46:
+   -------------------------------------------------
+    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:  }
+   -------------------------------------------------
+
+All arguments to "laplace_latent_solve" must be a data-only
+if used outside the generated quantities block. (Local variables are assumed
+to depend on parameters; same goes for function inputs unless they are marked
+with the keyword "data".)
+[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: