Expose builtin constraints

scripts/hooks/pre-commit

@@ -1,4 +1,5 @@
 #!/bin/bash
+set -e
 cd "$(git rev-parse --show-toplevel)" || exit
 dune runtest src test/unit && make format
 curl -X POST -F "jenkinsfile=<Jenkinsfile" "https://jenkins.flatironinstitute.org/pipeline-model-converter/validate" 2>/dev/null | tee /dev/tty | grep "Jenkinsfile successfully validated." -q

src/frontend/Typechecker.ml

@@ -449,12 +449,27 @@ let verify_fn_target_plus_equals cf loc id =
     else if in_lp_function cf || cf.current_block = Model then ()
     else Semantic_error.target_plusequals_outside_model_or_logprob loc |> error
 
-let verify_fn_jacobian_plus_equals cf loc id =
+let verify_fn_jacobian_plus_equals cf loc tenv id args =
   if
     String.is_suffix id.name ~suffix:"_jacobian"
-    && (not !Fun_kind.jacobian_compat_mode)
-    && not (in_jacobian_function cf || cf.current_block = TParam)
-  then Semantic_error.jacobian_plusequals_not_allowed loc |> error
+    && not !Fun_kind.jacobian_compat_mode
+  then
+    if not (in_jacobian_function cf || cf.current_block = TParam) then
+      Semantic_error.jacobian_plusequals_not_allowed loc |> error
+    else if
+      not
+        (List.exists args ~f:(fun e ->
+             UnsizedType.is_autodifftype e.emeta.ad_level))
+    then
+      let alt =
+        String.chop_suffix_exn ~suffix:"_jacobian" id.name ^ "_constrain" in
+      let message =
+        "Calling a _jacobian function without any parameter arguments still \
+         applies the Jacobian adjustments, ensure this is intentional!"
+        ^
+        if Env.mem tenv alt then " Consider using " ^ alt ^ " instead." else ""
+      in
+      warnings := (loc, message) :: !warnings
 
 (** Rng functions cannot be used in Tp or Model and only
     in function defs with the right suffix
@@ -703,7 +718,7 @@ and check_funapp loc cf tenv ~is_cond_dist id (es : Ast.typed_expression list) =
   verify_identifier id;
   name_check loc id;
   verify_fn_target_plus_equals cf loc id;
-  verify_fn_jacobian_plus_equals cf loc id;
+  verify_fn_jacobian_plus_equals cf loc tenv id es;
   verify_fn_rng cf loc id;
   verify_unnormalized cf loc id;
   res
@@ -950,7 +965,7 @@ let check_nr_fn_app loc cf tenv id es =
   let tes = List.map ~f:(check_expression cf tenv) es in
   verify_identifier id;
   verify_fn_target_plus_equals cf loc id;
-  verify_fn_jacobian_plus_equals cf loc id;
+  verify_fn_jacobian_plus_equals cf loc tenv id tes;
   verify_fn_rng cf loc id;
   check_nrfn loc tenv id tes
 

src/stan_math_backend/Lower_expr.ml

@@ -19,7 +19,23 @@ let fn_renames =
   @ [ ("lmultiply", "stan::math::multiply_log")
     ; ("lchoose", "stan::math::binomial_coefficient_log")
     ; ("std_normal_qf", "stan::math::inv_Phi")
-    ; ("integrate_ode", "stan::math::integrate_ode_rk45") ]
+    ; ("integrate_ode", "stan::math::integrate_ode_rk45")
+      (* constraints -- originally internal functions, may be worth renaming now *)
+    ; ("cholesky_factor_corr_jacobian", "stan::math::cholesky_corr_constrain")
+    ; ("cholesky_factor_cov_jacobian", "stan::math::cholesky_factor_constrain")
+    ; ("cholesky_factor_corr_constrain", "stan::math::cholesky_corr_constrain")
+    ; ("cholesky_factor_cov_constrain", "stan::math::cholesky_factor_constrain")
+    ; ("cholesky_factor_corr_unconstrain", "stan::math::cholesky_corr_free")
+    ; ("cholesky_factor_cov_unconstrain", "stan::math::cholesky_factor_free")
+    ; ("lower_bound_jacobian", "stan::math::lb_constrain")
+    ; ("upper_bound_jacobian", "stan::math::ub_constrain")
+    ; ("lower_upper_bound_jacobian", "stan::math::lub_constrain")
+    ; ("lower_bound_constrain", "stan::math::lb_constrain")
+    ; ("lower_bound_unconstrain", "stan::math::lb_free")
+    ; ("upper_bound_constrain", "stan::math::ub_constrain")
+    ; ("upper_bound_unconstrain", "stan::math::ub_free")
+    ; ("lower_upper_bound_constrain", "stan::math::lub_constrain")
+    ; ("lower_upper_bound_unconstrain", "stan::math::lub_free") ]
   |> String.Map.of_alist_exn
 
 let constraint_to_string = function
@@ -103,9 +119,11 @@ let promote_adtype =
       | _ -> accum)
     ~init:UnsizedType.DataOnly
 
-let suffix_args = function
+let suffix_args udf = function
   | Fun_kind.FnRng -> ["base_rng__"]
-  | FnTarget | FnJacobian -> ["lp__"; "lp_accum__"]
+  | FnTarget -> ["lp__"; "lp_accum__"]
+  | FnJacobian when udf -> ["lp__"; "lp_accum__"]
+  | FnJacobian -> ["lp__"]
   | FnPlain | FnLpdf _ | FnLpmf _ -> []
 
 let rec stantype_prim = function
@@ -118,7 +136,7 @@ let templates udf suffix =
   | Fun_kind.FnLpdf true | FnLpmf true -> [TemplateType "propto__"]
   | FnLpdf false | FnLpmf false -> [TemplateType "false"]
   | FnTarget when udf -> [TemplateType "propto__"]
-  | FnJacobian when udf -> [TemplateType "jacobian__"]
+  | FnJacobian -> [TemplateType "jacobian__"]
   | _ -> []
 
 let deserializer = Var "in__"
@@ -396,6 +414,12 @@ and lower_functionals fname suffix es mem_pattern =
 and lower_fun_app suffix fname es mem_pattern
     (ret_type : UnsizedType.returntype option) =
   let fname = Option.value (Map.find fn_renames fname) ~default:fname in
+  let fname =
+    (* Handle systematic renaming of math's constrain and free functions *)
+    match String.rsplit2 fname ~on:'_' with
+    | Some (f, "jacobian") -> f ^ "_constrain"
+    | Some (f, "unconstrain") -> f ^ "_free"
+    | _ -> fname in
   let special_options =
     [ Option.map ~f:lower_operator_app (Operator.of_string_opt fname)
     ; lower_misc_special_math_app fname mem_pattern ret_type
@@ -406,12 +430,12 @@ and lower_fun_app suffix fname es mem_pattern
   | None ->
       let fname = stan_namespace_qualify fname in
       let templates = templates false suffix in
-      let extras = suffix_args suffix |> List.map ~f:Exprs.to_var in
+      let extras = suffix_args false suffix |> List.map ~f:Exprs.to_var in
       Exprs.templated_fun_call fname templates (lower_exprs es @ extras)
 
 and lower_user_defined_fun f suffix es =
   let extra_args =
-    suffix_args suffix @ ["pstream__"] |> List.map ~f:Exprs.to_var in
+    suffix_args true suffix @ ["pstream__"] |> List.map ~f:Exprs.to_var in
   Exprs.templated_fun_call f (templates true suffix)
     ((lower_exprs ~promote_reals:true) es @ extra_args)
 

src/stan_math_signatures/Generate.ml

@@ -550,6 +550,36 @@ let add_binary_vec_int_int name supports_soa =
         , supports_soa ))
     (List.range 1 8)
 
+(* things like lower_bound require the second argument be a scalar or the same
+    - this is basically the general case above with the last nested loop removed
+*)
+let add_first_arg_vector_binary name supports_soa =
+  let vectors = [UnsizedType.UArray UReal; UVector; URowVector; UMatrix] in
+  add_unqualified (name, ReturnType UReal, [UReal; UReal], supports_soa);
+  List.iter
+    ~f:(fun i ->
+      List.iter
+        ~f:(fun j ->
+          let ty = bare_array_type (j, i) in
+          add_unqualified (name, ReturnType ty, [ty; ty], supports_soa);
+          add_unqualified (name, ReturnType ty, [ty; UReal], supports_soa))
+        vectors)
+    (List.range 0 8)
+
+let add_first_arg_vector_ternary name supports_soa =
+  let vectors = [UnsizedType.UArray UReal; UVector; URowVector; UMatrix] in
+  add_unqualified (name, ReturnType UReal, [UReal; UReal; UReal], supports_soa);
+  List.iter
+    ~f:(fun i ->
+      List.iter
+        ~f:(fun j ->
+          let ty = bare_array_type (j, i) in
+          add_unqualified (name, ReturnType ty, [ty; ty; ty], supports_soa);
+          add_unqualified (name, ReturnType ty, [ty; ty; UReal], supports_soa);
+          add_unqualified (name, ReturnType ty, [ty; UReal; ty], supports_soa))
+        vectors)
+    (List.range 0 8)
+
 let add_ternary name supports_soa =
   add_unqualified (name, ReturnType UReal, [UReal; UReal; UReal], supports_soa)
 
@@ -842,8 +872,35 @@ let () =
     , ReturnType UComplexRowVector
     , [UComplexRowVector; UComplex]
     , AoS );
+  List.iter
+    ~f:(fun i ->
+      let m = bare_array_type (UMatrix, i) in
+      let v = bare_array_type (UVector, i) in
+      add_unqualified ("corr_matrix_constrain", ReturnType m, [v; UInt], SoA);
+      add_unqualified ("corr_matrix_unconstrain", ReturnType v, [m], AoS);
+      add_unqualified ("corr_matrix_jacobian", ReturnType m, [v; UInt], SoA);
+      add_unqualified ("cov_matrix_constrain", ReturnType m, [v; UInt], SoA);
+      add_unqualified ("cov_matrix_unconstrain", ReturnType v, [m], AoS);
+      add_unqualified ("cov_matrix_jacobian", ReturnType m, [v; UInt], SoA))
+    (List.range 0 8);
   add_unqualified ("chol2inv", ReturnType UMatrix, [UMatrix], AoS);
   add_unqualified ("cholesky_decompose", ReturnType UMatrix, [UMatrix], SoA);
+  List.iter
+    ~f:(fun i ->
+      let m = bare_array_type (UMatrix, i) in
+      let v = bare_array_type (UVector, i) in
+      add_unqualified
+        ("cholesky_factor_corr_constrain", ReturnType m, [v; UInt], SoA);
+      add_unqualified
+        ("cholesky_factor_corr_unconstrain", ReturnType v, [m], AoS);
+      add_unqualified
+        ("cholesky_factor_corr_jacobian", ReturnType m, [v; UInt], SoA);
+      add_unqualified
+        ("cholesky_factor_cov_constrain", ReturnType m, [v; UInt; UInt], SoA);
+      add_unqualified ("cholesky_factor_cov_unconstrain", ReturnType v, [m], AoS);
+      add_unqualified
+        ("cholesky_factor_cov_jacobian", ReturnType m, [v; UInt; UInt], SoA))
+    (List.range 0 8);
   add_binary_vec_int_int "choose" AoS;
   add_unqualified ("col", ReturnType UVector, [UMatrix; UInt], AoS);
   add_unqualified ("col", ReturnType UComplexVector, [UComplexMatrix; UInt], SoA);
@@ -1542,6 +1599,12 @@ let () =
   add_unqualified ("logical_eq", ReturnType UInt, [UComplex; UComplex], SoA);
   add_unqualified ("logical_neq", ReturnType UInt, [UComplex; UReal], SoA);
   add_unqualified ("logical_neq", ReturnType UInt, [UComplex; UComplex], SoA);
+  add_first_arg_vector_binary "lower_bound_jacobian" SoA;
+  add_first_arg_vector_binary "lower_bound_constrain" SoA;
+  add_first_arg_vector_binary "lower_bound_unconstrain" AoS;
+  add_first_arg_vector_ternary "lower_upper_bound_jacobian" SoA;
+  add_first_arg_vector_ternary "lower_upper_bound_constrain" SoA;
+  add_first_arg_vector_ternary "lower_upper_bound_unconstrain" AoS;
   add_nullary "machine_precision";
   add_qualified
     ( "map_rect"
@@ -1834,6 +1897,9 @@ let () =
             ("num_elements", ReturnType UInt, [bare_array_type (t, i)], SoA))
         bare_types)
     (List.range 1 10);
+  add_first_arg_vector_ternary "offset_multiplier_jacobian" SoA;
+  add_first_arg_vector_ternary "offset_multiplier_constrain" SoA;
+  add_first_arg_vector_ternary "offset_multiplier_unconstrain" AoS;
   add_unqualified
     ("one_hot_int_array", ReturnType (UArray UInt), [UInt; UInt], SoA);
   add_unqualified ("one_hot_array", ReturnType (UArray UReal), [UInt; UInt], SoA);
@@ -1844,6 +1910,13 @@ let () =
   add_unqualified ("ones_array", ReturnType (UArray UReal), [UInt], SoA);
   add_unqualified ("ones_row_vector", ReturnType URowVector, [UInt], SoA);
   add_unqualified ("ones_vector", ReturnType UVector, [UInt], SoA);
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UVector, i) in
+      add_unqualified ("ordered_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("ordered_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("ordered_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
   add_unqualified
     ( "ordered_logistic_glm_lpmf"
     , ReturnType UReal
@@ -1914,6 +1987,13 @@ let () =
     , SoA );
   add_unqualified ("polar", ReturnType UComplex, [UReal; UReal], AoS);
   add_nullary "positive_infinity";
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UVector, i) in
+      add_unqualified ("positive_ordered_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("positive_ordered_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("positive_ordered_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
   add_binary_vec "pow" AoS;
   add_binary_vec_complex_complex "pow" AoS;
   add_unqualified ("prod", ReturnType UInt, [UArray UInt], AoS);
@@ -2133,6 +2213,13 @@ let () =
             , SoA ))
         (List.range 1 4))
     bare_types;
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UVector, i) in
+      add_unqualified ("simplex_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("simplex_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("simplex_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
   add_unqualified ("sin", ReturnType UComplex, [UComplex], AoS);
   add_unqualified ("sinh", ReturnType UComplex, [UComplex], AoS);
   add_unqualified ("singular_values", ReturnType UVector, [UMatrix], SoA);
@@ -2171,6 +2258,16 @@ let () =
   add_unqualified ("sort_indices_desc", ReturnType (UArray UInt), [UVector], AoS);
   add_unqualified
     ("sort_indices_desc", ReturnType (UArray UInt), [URowVector], AoS);
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UMatrix, i) in
+      add_unqualified ("stochastic_column_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("stochastic_column_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("stochastic_column_unconstrain", ReturnType t, [t], AoS);
+      add_unqualified ("stochastic_row_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("stochastic_row_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("stochastic_row_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
   add_unqualified ("squared_distance", ReturnType UReal, [UReal; UReal], SoA);
   add_unqualified ("squared_distance", ReturnType UReal, [UVector; UVector], SoA);
   add_unqualified
@@ -2231,6 +2328,13 @@ let () =
   add_unqualified ("sum", ReturnType UComplex, [UComplexVector], SoA);
   add_unqualified ("sum", ReturnType UComplex, [UComplexRowVector], SoA);
   add_unqualified ("sum", ReturnType UComplex, [UComplexMatrix], SoA);
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UVector, i) in
+      add_unqualified ("sum_to_zero_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("sum_to_zero_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("sum_to_zero_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
   (* TODO (future): SoA inside of tuples, update following signatures *)
   add_unqualified
     ("svd", ReturnType (UTuple [UMatrix; UVector; UMatrix]), [UMatrix], AoS);
@@ -2420,6 +2524,16 @@ let () =
     ("transpose", ReturnType UComplexVector, [UComplexRowVector], SoA);
   add_unqualified ("transpose", ReturnType UComplexMatrix, [UComplexMatrix], SoA);
   add_unqualified ("uniform_simplex", ReturnType UVector, [UInt], SoA);
+  List.iter
+    ~f:(fun i ->
+      let t = bare_array_type (UVector, i) in
+      add_unqualified ("unit_vector_jacobian", ReturnType t, [t], SoA);
+      add_unqualified ("unit_vector_constrain", ReturnType t, [t], SoA);
+      add_unqualified ("unit_vector_unconstrain", ReturnType t, [t], AoS))
+    (List.range 0 8);
+  add_first_arg_vector_binary "upper_bound_jacobian" SoA;
+  add_first_arg_vector_binary "upper_bound_constrain" SoA;
+  add_first_arg_vector_binary "upper_bound_unconstrain" AoS;
   add_unqualified ("variance", ReturnType UReal, [UArray UReal], SoA);
   add_unqualified ("variance", ReturnType UReal, [UVector], SoA);
   add_unqualified ("variance", ReturnType UReal, [URowVector], SoA);

test/integration/bad/err-jacobian-plusequals-scope3.stan

@@ -1,5 +1,5 @@
 functions {
-  real upper_bound_jacobian(real x, real ub) {
+  real my_upper_bound_jacobian(real x, real ub) {
     jacobian += x;
     return ub - exp(x);
   }
@@ -11,5 +11,5 @@ parameters {
 }
 
 model {
-  real b = upper_bound_jacobian(b_raw, ub);
+  real b = my_upper_bound_jacobian(b_raw, ub);
 }

test/integration/bad/err-jacobian-plusequals-scope5.stan

@@ -0,0 +1,9 @@
+
+parameters {
+  real b_raw;
+  real ub;
+}
+
+model {
+  real b = upper_bound_jacobian(b_raw, ub);
+}

test/integration/bad/stanc.expected

@@ -1074,11 +1074,11 @@ Semantic error in 'err-jacobian-plusequals-scope2.stan', line 3, column 4 to col
 The jacobian adjustment can only be applied in the transformed parameters block or in functions ending with _jacobian
 [exit 1]
   $ ../../../../install/default/bin/stanc err-jacobian-plusequals-scope3.stan
-Semantic error in 'err-jacobian-plusequals-scope3.stan', line 14, column 11 to column 42:
+Semantic error in 'err-jacobian-plusequals-scope3.stan', line 14, column 11 to column 45:
    -------------------------------------------------
     12:  
     13:  model {
-    14:    real b = upper_bound_jacobian(b_raw, ub);
+    14:    real b = my_upper_bound_jacobian(b_raw, ub);
                     ^
     15:  }
    -------------------------------------------------
@@ -1095,6 +1095,18 @@ Semantic error in 'err-jacobian-plusequals-scope4.stan', line 8, column 2 to col
      9:  }
    -------------------------------------------------
 
+The jacobian adjustment can only be applied in the transformed parameters block or in functions ending with _jacobian
+[exit 1]
+  $ ../../../../install/default/bin/stanc err-jacobian-plusequals-scope5.stan
+Semantic error in 'err-jacobian-plusequals-scope5.stan', line 8, column 11 to column 42:
+   -------------------------------------------------
+     6:  
+     7:  model {
+     8:    real b = upper_bound_jacobian(b_raw, ub);
+                    ^
+     9:  }
+   -------------------------------------------------
+
 The jacobian adjustment can only be applied in the transformed parameters block or in functions ending with _jacobian
 [exit 1]
   $ ../../../../install/default/bin/stanc err-minus-types.stan

test/integration/good/code-gen/builtin-constraint.stan

@@ -0,0 +1,18 @@
+data {
+  int<lower=1> K;
+}
+parameters {
+  vector[K-1] beta_raw;
+}
+transformed parameters {
+  sum_to_zero_vector[K] beta = sum_to_zero_jacobian(beta_raw);
+}
+model {
+  beta ~ normal(0, inv(sqrt(1 - inv(K))));
+}
+generated quantities {
+  vector[K-1] beta_recovered = sum_to_zero_unconstrain(beta);
+  if (max(abs(beta_recovered - beta_raw)) > 1e-10) {
+    fatal_error("beta_recovered does not match beta");
+  }
+}