Merge pull request #1516 from stan-dev/laplace-check-indexing-exprs

Laplace: also check into index expressions for incompatibilities.

Author: WardBrian

src/frontend/Typechecker.ml

@@ -635,6 +635,8 @@ let verify_second_order_derivative_compatibility (ast : typed_program) =
       let rec check_expr seen = function
         | {expr= FunApp (StanLib _, {name; _}, _); _}
           when Stan_math_signatures.lacks_higher_order_autodiff name ->
+            (* note: we could possibly check all the arguments are DataOnly
+               and still allow it, but those seem like mostly useless cases. *)
             Semantic_error.laplace_compatibility id_loc name |> error
         | {expr= FunApp (UserDefined _, name, es); _} ->
             (* we want the location to be the use-site no matter what *)
@@ -643,14 +645,15 @@ let verify_second_order_derivative_compatibility (ast : typed_program) =
         | {expr= Variable name; emeta= {type_= UFun _; _}} ->
             check_fun seen {name with id_loc}
         | e -> Ast.fold_expression check_expr fold_nop seen e.expr in
+      let check_lval acc l = fold_lval_with check_expr fold_nop acc l in
       let rec check_stmt seen s =
         match s.stmt with
         | NRFunApp (UserDefined _, name, es) ->
             let seen' = check_fun seen {name with id_loc} in
             List.fold ~f:check_expr ~init:seen' es
         | stmt ->
-            Ast.fold_statement check_expr check_stmt fold_nop fold_nop seen stmt
-      in
+            Ast.fold_statement check_expr check_stmt check_lval fold_nop seen
+              stmt in
       let visited' = Set.add visited fn_name in
       List.fold ~f:check_stmt ~init:visited' (get_function_bodies fn_name) in
   ignore

test/integration/bad/embedded_laplace/autodiff_incompatibility5.stan

@@ -0,0 +1,58 @@
+functions {
+
+  vector algebra_system(vector x, vector y, array[] real dat,
+                        array[] int dat_int) {
+    vector[2] f_x;
+    f_x[1] = x[1] - y[1];
+    f_x[2] = x[2] - y[2];
+    return f_x;
+  }
+
+
+  real helper(real eta){
+    array[3] real dat = {eta, 0.2, 0.3};
+    dat[cols(algebra_solver(algebra_system, [eta]', [1]', {0.1}, {2}))] = 0.0;
+
+    return dat[1];
+  }
+  // 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) + helper(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),
+                        theta_0,
+                        K_function, (x, n_obs, alpha, rho));
+}

test/integration/bad/embedded_laplace/stanc.expected

@@ -57,6 +57,21 @@ Semantic error in 'autodiff_incompatibility4.stan', line 51, column 43 to column
 The function 'reduce_sum', called by this likelihood function,
 does not currently support higher-order derivatives, and
 cannot be used in an embedded Laplace approximation.
+[exit 1]
+  $ ../../../../../install/default/bin/stanc autodiff_incompatibility5.stan
+Semantic error in 'autodiff_incompatibility5.stan', line 55, column 43 to column 54:
+   -------------------------------------------------
+    53:  
+    54:  generated quantities {
+    55:    vector[n_obs] theta = laplace_latent_rng(ll_function, (eta, log_ye, y),
+                                                    ^
+    56:                          theta_0,
+    57:                          K_function, (x, n_obs, alpha, rho));
+   -------------------------------------------------
+
+The function 'algebra_solver', called by this likelihood function,
+does not currently support higher-order derivatives, and
+cannot be used in an embedded Laplace approximation.
 [exit 1]
   $ ../../../../../install/default/bin/stanc bad_callback1.stan
 Semantic error in 'bad_callback1.stan', line 37, column 29 to column 40: