Merge pull request #1046 from WardBrian/lpdf_typing

User-defined densities allowed over nested and complex values.

Author: WardBrian

src/frontend/Deprecation_analysis.ml

@@ -82,9 +82,8 @@ let update_suffix name type_ =
   let open String in
   if is_suffix ~suffix:"_cdf_log" name then drop_suffix name 8 ^ "_lcdf"
   else if is_suffix ~suffix:"_ccdf_log" name then drop_suffix name 9 ^ "_lccdf"
-  else if Middle.UnsizedType.is_real_type type_ then
-    drop_suffix name 4 ^ "_lpdf"
-  else drop_suffix name 4 ^ "_lpmf"
+  else if Middle.UnsizedType.is_int_type type_ then drop_suffix name 4 ^ "_lpmf"
+  else drop_suffix name 4 ^ "_lpdf"
 
 let find_udf_log_suffix = function
   | { stmt=

src/frontend/Typechecker.ml

@@ -59,7 +59,7 @@ let context block =
 let calculate_autodifftype cf origin ut =
   match origin with
   | Env.(Param | TParam | Model | Functions)
-    when not (UnsizedType.contains_int ut || cf.current_block = GQuant) ->
+    when not (UnsizedType.is_int_type ut || cf.current_block = GQuant) ->
       UnsizedType.AutoDiffable
   | _ -> DataOnly
 
@@ -1137,7 +1137,7 @@ and verify_transformed_param_ty loc cf is_global unsized_ty =
   if
     is_global
     && (cf.current_block = Param || cf.current_block = TParam)
-    && UnsizedType.contains_int unsized_ty
+    && UnsizedType.is_int_type unsized_ty
   then Semantic_error.transformed_params_int loc |> error
 
 and check_sizedtype cf tenv sizedty =
@@ -1268,7 +1268,7 @@ and verify_pdf_fundef_first_arg_ty loc id arg_tys =
   if String.is_suffix id.name ~suffix:"_lpdf" then
     let rt = List.hd arg_tys |> Option.map ~f:snd in
     match rt with
-    | Some rt when UnsizedType.is_real_type rt -> ()
+    | Some rt when not (UnsizedType.is_int_type rt) -> ()
     | _ -> Semantic_error.prob_density_non_real_variate loc rt |> error
 
 and verify_pmf_fundef_first_arg_ty loc id arg_tys =

src/middle/UnsizedType.ml

@@ -130,13 +130,9 @@ let rec common_type = function
   | _, _ -> None
 
 (* -- Helpers -- *)
-let is_real_type = function
-  | UReal | UVector | URowVector | UMatrix
-   |UArray UReal
-   |UArray UVector
-   |UArray URowVector
-   |UArray UMatrix ->
-      true
+let rec is_real_type = function
+  | UReal | UVector | URowVector | UMatrix -> true
+  | UArray x -> is_real_type x
   | _ -> false
 
 let rec is_autodiffable = function
@@ -145,17 +141,15 @@ let rec is_autodiffable = function
   | _ -> false
 
 let is_scalar_type = function UReal | UInt -> true | _ -> false
-let is_int_type = function UInt | UArray UInt -> true | _ -> false
+
+let rec is_int_type ut =
+  match ut with UInt -> true | UArray ut -> is_int_type ut | _ -> false
 
 let is_eigen_type ut =
   match ut with UVector | URowVector | UMatrix -> true | _ -> false
 
 let is_fun_type = function UFun _ | UMathLibraryFunction -> true | _ -> false
 
-(** Detect if type contains an integer *)
-let rec contains_int ut =
-  match ut with UInt -> true | UArray ut -> contains_int ut | _ -> false
-
 let rec is_indexing_matrix = function
   | UArray t, _ :: idcs -> is_indexing_matrix (t, idcs)
   | UMatrix, [] -> false

src/stan_math_backend/Stan_math_code_gen.ml

@@ -56,7 +56,7 @@ let pp_located ppf _ =
 (** Detect if argument requires C++ template *)
 let arg_needs_template = function
   | UnsizedType.DataOnly, _, t -> UnsizedType.is_eigen_type t
-  | _, _, t when UnsizedType.contains_int t -> false
+  | _, _, t when UnsizedType.is_int_type t -> false
   | _ -> true
 
 (** Print template arguments for C++ functions that need templates
@@ -107,7 +107,7 @@ let pp_promoted_scalar ppf args =
 let pp_returntype ppf arg_types rt =
   let scalar = str "%a" pp_promoted_scalar arg_types in
   match rt with
-  | Some ut when UnsizedType.contains_int ut ->
+  | Some ut when UnsizedType.is_int_type ut ->
       pf ppf "%a@," pp_unsizedtype_custom_scalar ("int", ut)
   | Some ut -> pf ppf "%a@," pp_unsizedtype_custom_scalar (scalar, ut)
   | None -> pf ppf "void@,"

test/integration/good/fun-defs-lpdf.stan

@@ -2,10 +2,20 @@ functions {
   real bar_baz_lpdf(real a, real b) {
     return a / b;
   }
+  real foo_bar_lpdf(array[,,,] real x){
+    return 1.0;
+  }
+  real baz_foo_lpdf(complex z, real a){
+    return get_imag(z) * a;
+  }
 }
 parameters {
   real y;
+  complex z;
+  array[1,1,1,1] real arr;
 }
 model {
   y ~ bar_baz(3.2);
+  z ~ baz_foo(1.5);
+  arr ~ foo_bar();
 }

test/integration/good/pretty.expected

@@ -3097,12 +3097,22 @@ functions {
   real bar_baz_lpdf(real a, real b) {
     return a / b;
   }
+  real foo_bar_lpdf(array[,,,] real x) {
+    return 1.0;
+  }
+  real baz_foo_lpdf(complex z, real a) {
+    return get_imag(z) * a;
+  }
 }
 parameters {
   real y;
+  complex z;
+  array[1, 1, 1, 1] real arr;
 }
 model {
   y ~ bar_baz(3.2);
+  z ~ baz_foo(1.5);
+  arr ~ foo_bar();
 }
 
   $ ../../../../install/default/bin/stanc --auto-format fun-return-typ1.stan