Merge pull request #1586 from stan-dev/laplace-tol-updates

Update laplace _tol functions for allow_fallthrough argument, tuple of control params

Author: WardBrian

src/frontend/Semantic_error.ml

@@ -77,13 +77,16 @@ module TypeError = struct
 
   let laplace_tolerance_arg_name n =
     match n with
-    | 1 -> "first control parameter (initial guess)"
-    | 2 -> "second control parameter (tolerance)"
-    | 3 -> "third control parameter (max_num_steps)"
-    | 4 -> "fourth control parameter (hessian_block_size)"
-    | 5 -> "fifth control parameter (solver)"
-    | 6 -> "sixth control parameter (max_steps_line_search)"
-    | n -> Fmt.str "%a control parameter" (Fmt.ordinal ()) n
+    | 1 -> "first element of the control parameter tuple (initial guess)"
+    | 2 -> "second element of the control parameter tuple (tolerance)"
+    | 3 -> "third element of the control parameter tuple (max_num_steps)"
+    | 4 -> "fourth element of the control parameter tuple (hessian_block_size)"
+    | 5 -> "fifth element of the control parameter tuple (solver)"
+    | 6 ->
+        "sixth element of the control parameter tuple (max_steps_line_search)"
+    | 7 -> "seventh element of the control parameter tuple (allow_fallthrough)"
+    | n ->
+        Fmt.str "%a element of the control parameter tuple" (Fmt.ordinal ()) n
 
   let rec expected_types : UnsizedType.t Common.Nonempty_list.t Fmt.t =
     let ust = expected_style UnsizedType.pp in
@@ -230,17 +233,23 @@ module TypeError = struct
            an embedded Laplace approximation."
           quoted banned_function
     | IlltypedLaplaceTooMany (name, n_args) ->
-        Fmt.pf ppf
-          "Received %d extra %a at the end of the call to %a.@ Did you mean to \
-           call the _tol version?"
+        Fmt.pf ppf "Received %d extra %a at the end of the call to %a.@ %s"
           n_args arguments n_args quoted name
-    (* For tolerances, because these come at the end, we want to update their
-       position number accordingly, which is why these reimplement some of the
-       printing from [SignatureMismatch] *)
+          (if String.is_substring ~substring:"_tol" name then
+             "Only a single tuple of control parameters is expected."
+           else if n_args = 1 then "Did you mean to call the _tol version?"
+           else "Did you mean to call the _tol version with a tuple of these?")
+    | IlltypedLaplaceTolArgs (name, ArgNumMismatch (_, 0)) ->
+        Fmt.pf ppf
+          "Missing control parameter tuple at the end of the call to %a.@ \
+           Expected a tuple of %a arguments for the control parameters."
+          quoted name (expected_style Fmt.int)
+          (List.length Stan_math_signatures.laplace_tolerance_argument_types)
     | IlltypedLaplaceTolArgs (name, ArgNumMismatch (_, found)) ->
         Fmt.pf ppf
-          "@[<v>Received %a control %a at the end of the call to %a.@ Expected \
-           %a arguments for the control parameters instead.@]"
+          "@[<v>Received a tuple of %a control %a at the end of the call to \
+           %a.@ Expected tuple of %a arguments for the control parameters \
+           instead.@]"
           (actual_style Fmt.int) found arguments found quoted name
           (expected_style Fmt.int)
           (List.length Stan_math_signatures.laplace_tolerance_argument_types)

src/frontend/Typechecker.ml

@@ -716,21 +716,33 @@ let check_function_callable_with_tuple cf tenv caller_id fname
 let verify_laplace_control_args loc id (args : typed_expression list) =
   match (String.is_substring ~substring:"_tol" id.name, args) with
   | false, [] -> ()
-  | true, _ -> (
-      let arg_tys = List.map ~f:arg_type args in
+  | true, [arg] -> (
+      let arg_tys =
+        check_texpression_is_tuple arg
+          ("Control arguments for '" ^ id.name ^ "'") in
       match
         SignatureMismatch.check_compatible_arguments_mod_conv
           Stan_math_signatures.laplace_tolerance_argument_types arg_tys
       with
       | Ok _ -> ()
       | Error err ->
           let loc =
-            let which_arg = match err with ArgError (i, _) -> i | _ -> 0 in
-            List.nth args which_arg
-            |> Option.value_map ~f:(fun expr -> expr.emeta.loc) ~default:loc
-          in
+            let which_arg = match err with ArgError (i, _) -> i - 1 | _ -> 0 in
+            let elts = match arg.expr with TupleExpr elts -> elts | _ -> [] in
+            List.nth elts which_arg
+            |> Option.value_map ~f:(fun e -> e.emeta.loc) ~default:loc in
           Semantic_error.illtyped_laplace_tolerance_args loc id.name err
           |> error)
+  | true, [] ->
+      Semantic_error.illtyped_laplace_tolerance_args loc id.name
+        (SignatureMismatch.check_compatible_arguments_mod_conv
+           Stan_math_signatures.laplace_tolerance_argument_types []
+        |> Result.error |> Option.value_exn)
+      |> error
+  | true, _ :: a :: _ ->
+      Semantic_error.illtyped_laplace_extra_args a.emeta.loc id.name
+        (List.length args - 1)
+      |> error
   | false, a :: _ ->
       Semantic_error.illtyped_laplace_extra_args a.emeta.loc id.name
         (List.length args)

src/stan_math_signatures/Stan_math_signatures.ml

@@ -169,14 +169,7 @@ let embedded_laplace_functions =
   |> String.Set.of_list
 
 let is_embedded_laplace_fn name =
-  (* TEMPORARY: remove after https://github.com/stan-dev/math/pull/3250 or
-     similar is merged *)
-  (not
-     (let version = "%%VERSION%%" in
-      String.equal "v2.38.0" version
-      || String.equal "v2.38.0-rc"
-           (String.sub ~pos:0 ~len:(String.length version - 1) version)))
-  && Set.mem embedded_laplace_functions (Utils.stdlib_distribution_name name)
+  Set.mem embedded_laplace_functions (Utils.stdlib_distribution_name name)
 
 let laplace_helper_lik_args =
   [ ( "bernoulli_logit"
@@ -206,7 +199,8 @@ let laplace_tolerance_argument_types =
     [ (AutoDiffable, UVector) (* theta_0 *); (DataOnly, UReal) (* tolerance *)
     ; (DataOnly, UInt) (* max_num_steps *)
     ; (DataOnly, UInt) (* hessian_block_size *); (DataOnly, UInt) (* solver *)
-    ; (DataOnly, UInt) (* max_steps_line_search *) ]
+    ; (DataOnly, UInt) (* max_steps_line_search *)
+    ; (DataOnly, UInt) (* allow_fallthrough *) ]
 
 let is_special_function_name name =
   is_stan_math_variadic_function_name name

test/integration/bad/embedded_laplace/bad_theta0.stan

@@ -30,7 +30,7 @@ transformed data {
 
   // control parameters for Laplace approximation
   real tolerance = 1e-6;
-  int max_num_steps = 100, hessian_block_size = 1, solver = 1, max_steps_line_search = 0;
+  int max_num_steps = 100, hessian_block_size = 1, solver = 1, max_steps_line_search = 0, allow_fallthrough = 1;
 }
 parameters {
   real<lower=0> alpha;
@@ -39,6 +39,7 @@ parameters {
 }
 model {
   target += laplace_marginal_tol(ll_function, (eta, log_ye, y),
-                                  K_function, (x, n_obs, alpha, rho), theta_0, tolerance, max_num_steps,
-              hessian_block_size, solver, max_steps_line_search);
+                                  K_function, (x, n_obs, alpha, rho),
+        (theta_0, tolerance, max_num_steps, hessian_block_size,
+        solver, max_steps_line_search, allow_fallthrough));
 }

test/integration/bad/embedded_laplace/bad_tol0.stan

@@ -0,0 +1,41 @@
+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);
+  vector[n_obs] theta_0 = rep_vector(0.0, n_obs); // initial guess
+}
+parameters {
+  real<lower=0> alpha;
+  real<lower=0> rho;
+  real<lower=0> eta;
+}
+
+
+generated quantities {
+  vector[n_obs] theta = laplace_latent_rng(ll_function, (eta, log_ye, y),
+                    K_function, (x, n_obs, alpha, rho), theta_0);
+}

test/integration/bad/embedded_laplace/bad_tol10.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);
+  vector[n_obs] theta_0 = rep_vector(0.0, n_obs); // initial guess
+}
+parameters {
+  real<lower=0> alpha;
+  real<lower=0> rho;
+  real<lower=0> eta;
+}
+
+generated quantities {
+  vector[n_obs] theta = laplace_latent_tol_rng(ll_function, (eta, log_ye, y),
+                    K_function, (x, n_obs, alpha, rho), (theta_0, 1,2,3,4,5,0.1));
+}

test/integration/bad/embedded_laplace/bad_tol11.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);
+  vector[n_obs] theta_0 = rep_vector(0.0, n_obs); // initial guess
+}
+parameters {
+  real<lower=0> alpha;
+  real<lower=0> rho;
+  real<lower=0> eta;
+}
+
+generated quantities {
+  vector[n_obs] theta = laplace_latent_tol_rng(ll_function, (eta, log_ye, y),
+                    K_function, (x, n_obs, alpha, rho), (theta_0, 1,2,3,4,5,0),1);
+}

test/integration/bad/embedded_laplace/bad_tol3.stan

@@ -36,5 +36,5 @@ parameters {
 
 generated quantities {
   vector[n_obs] theta = laplace_latent_tol_rng(ll_function, (eta, log_ye, y),
-                    K_function, (x, n_obs, alpha, rho), theta_0, 1,2,3,4);
+                    K_function, (x, n_obs, alpha, rho), (theta_0, 1,2,3,4));
 }

test/integration/bad/embedded_laplace/bad_tol4.stan

@@ -34,6 +34,7 @@ transformed data {
   int hessian_block_size = 1;
   int solver = 1;
   int max_steps_line_search = 0;
+  int allow_fallthrough = 1;
 }
 parameters {
   real<lower=0> alpha;
@@ -42,7 +43,7 @@ parameters {
 }
 model {
   target += laplace_marginal_tol(ll_function, (eta, log_ye, y),
-                                K_function, (x, n_obs, alpha, rho), theta_0,
+                                K_function, (x, n_obs, alpha, rho), (theta_0,
                                 eta, max_num_steps, hessian_block_size,
-                                solver, max_steps_line_search);
+                                solver, max_steps_line_search, allow_fallthrough));
 }

test/integration/bad/embedded_laplace/bad_tol5.stan

@@ -36,5 +36,5 @@ parameters {
 
 generated quantities {
   vector[n_obs] theta = laplace_latent_tol_rng(ll_function, (eta, log_ye, y),
-                    K_function, (x, n_obs, alpha, rho), theta_0, 1,2,3,4,5.5);
+                    K_function, (x, n_obs, alpha, rho), (theta_0, 1,2,3,4,5.5,0));
 }