Promotions from signatures

Author: WardBrian

src/frontend/Ast.ml

@@ -251,9 +251,14 @@ type typed_program = typed_statement program [@@deriving sexp, compare, map]
 (** Forgetful function from typed to untyped expressions *)
 let rec untyped_expression_of_typed_expression ({expr; emeta} : typed_expression)
     : untyped_expression =
-  { expr=
-      map_expression untyped_expression_of_typed_expression (fun _ -> ()) expr
-  ; emeta= {loc= emeta.loc} }
+  match expr with
+  | Promotion (e, _) -> untyped_expression_of_typed_expression e
+  | _ ->
+      { expr=
+          map_expression untyped_expression_of_typed_expression
+            (fun _ -> ())
+            expr
+      ; emeta= {loc= emeta.loc} }
 
 let rec untyped_lvalue_of_typed_lvalue ({lval; lmeta} : typed_lval) :
     untyped_lval =

src/frontend/SignatureMismatch.ml

@@ -110,41 +110,52 @@ let rec compare_errors e1 e2 =
       | SuffixMismatch _, _ | _, InputMismatch _ -> -1
       | InputMismatch _, _ | _, SuffixMismatch _ -> 1 ) )
 
+type promotions = None | RealPromotion | ComplexPromotion
+
+(*
+PROMOTION TODO: Not sure what depth actually means, or how it needs to change
+*)
 let rec check_same_type depth t1 t2 =
   let wrap_func = Result.map_error ~f:(fun e -> TypeMismatch (t1, t2, Some e)) in
   match (t1, t2) with
-  | t1, t2 when t1 = t2 -> Ok ()
-  | UnsizedType.(UReal, UInt) when depth < 1 -> Ok ()
-  | UnsizedType.(UComplex, UInt) when depth < 1 -> Ok ()
-  | UnsizedType.(UComplex, UReal) when depth < 1 -> Ok ()
+  | t1, t2 when t1 = t2 -> Ok None
+  | UnsizedType.(UReal, UInt) when depth < 1 -> Ok RealPromotion
+  | UnsizedType.(UComplex, UInt) when depth < 1 -> Ok ComplexPromotion
+  | UnsizedType.(UComplex, UReal) when depth < 1 -> Ok ComplexPromotion
   | UFun (_, _, s1, _), UFun (_, _, s2, _)
     when Fun_kind.without_propto s1 <> Fun_kind.without_propto s2 ->
       Error
         (SuffixMismatch (Fun_kind.without_propto s1, Fun_kind.without_propto s2))
       |> wrap_func
   | UFun (_, rt1, _, _), UFun (_, rt2, _, _) when rt1 <> rt2 ->
       Error (ReturnTypeMismatch (rt1, rt2)) |> wrap_func
-  | UFun (l1, _, _, _), UFun (l2, _, _, _) ->
-      check_compatible_arguments (depth + 1) l2 l1
-      |> Result.map_error ~f:(fun e -> InputMismatch e)
-      |> wrap_func
+  | UFun (l1, _, _, _), UFun (l2, _, _, _) -> (
+    match check_compatible_arguments (depth + 1) l2 l1 with
+    | Ok _ -> Ok None
+    | Error e -> Error (InputMismatch e) |> wrap_func )
+  (* PROMOTION TODO: Need to allow array promotion, this doesn't quite work? *)
+  | UArray ut1, UArray ut2 ->
+      check_same_type depth ut1 ut2
+      |> Result.map_error ~f:(function
+           | TypeMismatch (_, _, _) -> TypeMismatch (t1, t2, None)
+           | e -> e )
   | t1, t2 -> Error (TypeMismatch (t1, t2, None))
 
-and check_compatible_arguments depth args1 args2 =
-  match List.zip args1 args2 with
+and check_compatible_arguments depth typs args2 :
+    (promotions list, function_mismatch) result =
+  match List.zip typs args2 with
   | List.Or_unequal_lengths.Unequal_lengths ->
-      Error (ArgNumMismatch (List.length args1, List.length args2))
+      Error (ArgNumMismatch (List.length typs, List.length args2))
   | Ok l ->
-      List.find_mapi l ~f:(fun i ((ad1, ut1), (ad2, ut2)) ->
+      List.mapi l ~f:(fun i ((ad1, ut1), (ad2, ut2)) ->
           match check_same_type depth ut1 ut2 with
-          | Error e -> Some (ArgError (i + 1, e))
-          | Ok _ ->
-              if ad1 = ad2 then None
+          | Error e -> Error (ArgError (i + 1, e))
+          | Ok p ->
+              if ad1 = ad2 then Ok p
               else if depth < 2 && UnsizedType.autodifftype_can_convert ad1 ad2
-              then None
-              else Some (ArgError (i + 1, DataOnlyError)) )
-      |> Option.map ~f:Result.fail
-      |> Option.value ~default:(Ok ())
+              then Ok p
+              else Error (ArgError (i + 1, DataOnlyError)) )
+      |> Result.all
 
 let check_compatible_arguments_mod_conv = check_compatible_arguments 0
 let max_n_errors = 5
@@ -157,13 +168,19 @@ let extract_function_types f =
       Some (return, args, (fun x -> UserDefined x), mem)
   | _ -> None
 
-let arg_type x = Ast.(x.emeta.ad_level, x.emeta.type_)
-let get_arg_types = List.map ~f:arg_type
+let promote es promotions =
+  List.map2_exn es promotions ~f:(fun (exp : Ast.typed_expression) prom ->
+      let emeta = exp.emeta in
+      match prom with
+      | RealPromotion when emeta.type_ <> UReal ->
+          Ast.{expr= Ast.Promotion (exp, UReal); emeta}
+      | ComplexPromotion when emeta.type_ <> UComplex ->
+          {expr= Promotion (exp, UComplex); emeta}
+      | _ -> exp )
 
-let returntype env name arg_exprs =
+let returntype env name args =
   (* NB: Variadic arguments are special-cased in the typechecker and not handled here *)
   let name = Utils.stdlib_distribution_name name in
-  let args = get_arg_types arg_exprs in
   Environment.find env name
   |> List.filter_map ~f:extract_function_types
   |> List.sort ~compare:(fun (x, _, _, _) (y, _, _, _) ->
@@ -173,7 +190,7 @@ let returntype env name arg_exprs =
        ~f:(fun errors (rt, tys, funkind_constructor, _) ->
          match check_compatible_arguments 0 tys args with
          (* TODO instead of unit, return Ast.typed_expr list which could contain promotions*)
-         | Ok () -> Stop (Ok (rt, funkind_constructor))
+         | Ok p -> Stop (Ok (rt, funkind_constructor, p))
          | Error e -> Continue (((rt, tys), e) :: errors) )
        ~finish:(fun errors ->
          let errors =
@@ -203,13 +220,13 @@ let check_variadic_args allow_lpdf mandatory_arg_tys mandatory_fun_arg_tys
       if suffix = FnPlain || (allow_lpdf && suffix = FnLpdf ()) then
         match check_compatible_arguments 1 mandatory mandatory_fun_arg_tys with
         | Error x -> wrap_func_error (InputMismatch x)
-        | Ok () -> (
+        | Ok _ -> (
           match check_same_type 1 return_type fun_return with
           | Error _ ->
               wrap_func_error
                 (ReturnTypeMismatch
                    (ReturnType fun_return, ReturnType return_type) )
-          | Ok () ->
+          | Ok _ ->
               let expected_args =
                 ((UnsizedType.AutoDiffable, func_type) :: mandatory_arg_tys)
                 @ variadic_arg_tys in

src/frontend/SignatureMismatch.mli

@@ -19,16 +19,23 @@ type signature_error =
   (UnsizedType.returntype * (UnsizedType.autodifftype * UnsizedType.t) list)
   * function_mismatch
 
+type promotions = private None | RealPromotion | ComplexPromotion
+
 val check_compatible_arguments_mod_conv :
      (UnsizedType.autodifftype * UnsizedType.t) list
   -> (UnsizedType.autodifftype * UnsizedType.t) list
-  -> (unit (* Ast.typed_expression list *), function_mismatch) result
+  -> (promotions list, function_mismatch) result
+
+val promote :
+  Ast.typed_expression list -> promotions list -> Ast.typed_expression list
 
 val returntype :
      Environment.t
   -> string
-  -> Ast.typed_expression list
-  -> ( UnsizedType.returntype * (bool Middle.Fun_kind.suffix -> Ast.fun_kind)
+  -> (UnsizedType.autodifftype * UnsizedType.t) list
+  -> ( UnsizedType.returntype
+       * (bool Middle.Fun_kind.suffix -> Ast.fun_kind)
+       * promotions list
      , signature_error list * bool )
      result
 
@@ -38,7 +45,7 @@ val check_variadic_args :
   -> (UnsizedType.autodifftype * UnsizedType.t) list
   -> UnsizedType.t
   -> (UnsizedType.autodifftype * UnsizedType.t) list
-  -> ( unit
+  -> ( promotions list
      , (UnsizedType.autodifftype * UnsizedType.t) list * function_mismatch )
      result
 

src/frontend/Typechecker.ml

@@ -424,15 +424,17 @@ let check_fn ~is_cond_dist loc tenv id es =
             (Env.nearest_ident tenv id.name) )
       |> error
   | _ (* a function *) -> (
-    match SignatureMismatch.returntype tenv id.name es with
-    | Ok (Void, _) ->
+    match SignatureMismatch.returntype tenv id.name (get_arg_types es) with
+    | Ok (Void, _, _) ->
         Semantic_error.returning_fn_expected_nonreturning_found loc id.name
         |> error
-    | Ok (ReturnType ut, fnk) ->
+    | Ok (ReturnType ut, fnk, promotions) ->
         mk_typed_expression
           ~expr:
             (mk_fun_app ~is_cond_dist
-               (fnk (Fun_kind.suffix_from_name id.name), id, es) )
+               ( fnk (Fun_kind.suffix_from_name id.name)
+               , id
+               , SignatureMismatch.promote es promotions ) )
           ~ad_level:(expr_ad_lub es) ~type_:ut ~loc
     | Error x ->
         es
@@ -458,9 +460,11 @@ let check_reduce_sum ~is_cond_dist loc id es =
         SignatureMismatch.check_variadic_args true mandatory_args
           mandatory_fun_args UReal (get_arg_types es)
       with
-      | Ok () ->
+      | Ok promotions ->
           mk_typed_expression
-            ~expr:(mk_fun_app ~is_cond_dist (StanLib FnPlain, id, es))
+            ~expr:
+              (mk_fun_app ~is_cond_dist
+                 (StanLib FnPlain, id, SignatureMismatch.promote es promotions) )
             ~ad_level:(expr_ad_lub es) ~type_:UnsizedType.UReal ~loc
       | Error (expected_args, err) ->
           Semantic_error.illtyped_reduce_sum loc id.name
@@ -498,9 +502,11 @@ let check_variadic_ode ~is_cond_dist loc id es =
       Stan_math_signatures.variadic_ode_mandatory_fun_args
       Stan_math_signatures.variadic_ode_fun_return_type (get_arg_types es)
   with
-  | Ok () ->
+  | Ok promotions ->
       mk_typed_expression
-        ~expr:(mk_fun_app ~is_cond_dist (StanLib FnPlain, id, es))
+        ~expr:
+          (mk_fun_app ~is_cond_dist
+             (StanLib FnPlain, id, SignatureMismatch.promote es promotions) )
         ~ad_level:(expr_ad_lub es)
         ~type_:Stan_math_signatures.variadic_ode_return_type ~loc
   | Error (expected_args, err) ->
@@ -702,12 +708,16 @@ let check_nrfn loc tenv id es =
         (Env.nearest_ident tenv id.name)
       |> error
   | _ (* a function *) -> (
-    match SignatureMismatch.returntype tenv id.name es with
-    | Ok (Void, fnk) ->
+    match SignatureMismatch.returntype tenv id.name (get_arg_types es) with
+    | Ok (Void, fnk, promotions) ->
         mk_typed_statement
-          ~stmt:(NRFunApp (fnk (Fun_kind.suffix_from_name id.name), id, es))
+          ~stmt:
+            (NRFunApp
+               ( fnk (Fun_kind.suffix_from_name id.name)
+               , id
+               , SignatureMismatch.promote es promotions ) )
           ~return_type:NoReturnType ~loc
-    | Ok (ReturnType _, _) ->
+    | Ok (ReturnType _, _, _) ->
         Semantic_error.nonreturning_fn_expected_returning_found loc id.name
         |> error
     | Error x ->

src/stan_math_backend/Expression_gen.ml

@@ -528,8 +528,10 @@ and pp_expr ppf Expr.Fixed.({pattern; meta} as e) =
   | Lit (Imaginary, s) -> pf ppf "to_complex(0, %s)" s
   | Lit ((Real | Int), s) -> pf ppf "%s" s
   | Promotion (expr, ut) ->
-      pf ppf "stan::math::promote_scalar_t<%a>(%a)" pp_unsizedtype_local
-        (meta.adlevel, ut) pp_expr expr
+      if is_scalar expr && ut = UReal then pp_expr ppf expr
+      else
+        pf ppf "stan::math::promote_scalar_t<%a>(%a)" pp_unsizedtype_local
+          (meta.adlevel, ut) pp_expr expr
   | FunApp
       ( StanLib (op, _, _)
       , [ { meta= {type_= URowVector; _}

test/integration/bad/algebra_solver/stanc.expected

@@ -128,47 +128,9 @@ where F2 = (vector, vector, data array[] real, data array[] int) => vector
 (<F2>, vector, vector, data array[] real, data array[] int) => vector
   Expected 5 arguments but found 8 arguments.
   $ ../../../../../install/default/bin/stanc bad_newton_x_r_type.stan
-Semantic error in 'bad_newton_x_r_type.stan', line 31, column 10 to column 69:
-   -------------------------------------------------
-    29:  transformed parameters {
-    30:    vector[2] y_s_p;
-    31:    y_s_p = algebra_solver_newton(algebra_system, y, theta_p, x_r, x_i);
-                   ^
-    32:  }
-    33:  
-   -------------------------------------------------
 
-Ill-typed arguments supplied to function 'algebra_solver_newton':
-(<F1>, vector, vector, array[] int, array[] int)
-where F1 = (vector, vector, array[] real, array[] int) => vector
-Available signatures:
-(<F2>, vector, vector, data array[] real, data array[] int) => vector
-where F2 = (vector, vector, data array[] real, data array[] int) => vector
-  The fourth argument must be array[] real but got array[] int
-(<F2>, vector, vector, data array[] real, data array[] int, data real,
- data real, data real) => vector
-  Expected 8 arguments but found 5 arguments.
   $ ../../../../../install/default/bin/stanc bad_newton_x_r_type_control.stan
-Semantic error in 'bad_newton_x_r_type_control.stan', line 31, column 10 to column 85:
-   -------------------------------------------------
-    29:  transformed parameters {
-    30:    vector[2] y_s_p;
-    31:    y_s_p = algebra_solver_newton(algebra_system, y, theta_p, x_r, x_i, 0.01, 0.01, 10);
-                   ^
-    32:  }
-    33:  
-   -------------------------------------------------
 
-Ill-typed arguments supplied to function 'algebra_solver_newton':
-(<F1>, vector, vector, array[] int, array[] int, real, real, int)
-where F1 = (vector, vector, array[] real, array[] int) => vector
-Available signatures:
-(<F2>, vector, vector, data array[] real, data array[] int, data real,
- data real, data real) => vector
-where F2 = (vector, vector, data array[] real, data array[] int) => vector
-  The fourth argument must be array[] real but got array[] int
-(<F2>, vector, vector, data array[] real, data array[] int) => vector
-  Expected 5 arguments but found 8 arguments.
   $ ../../../../../install/default/bin/stanc bad_newton_x_r_var_type.stan
 Semantic error in 'bad_newton_x_r_var_type.stan', line 31, column 10 to column 71:
    -------------------------------------------------
@@ -374,47 +336,9 @@ where F2 = (vector, vector, data array[] real, data array[] int) => vector
 (<F2>, vector, vector, data array[] real, data array[] int) => vector
   Expected 5 arguments but found 8 arguments.
   $ ../../../../../install/default/bin/stanc bad_x_r_type.stan
-Semantic error in 'bad_x_r_type.stan', line 31, column 10 to column 62:
-   -------------------------------------------------
-    29:  transformed parameters {
-    30:    vector[2] y_s_p;
-    31:    y_s_p = algebra_solver(algebra_system, y, theta_p, x_r, x_i);
-                   ^
-    32:  }
-    33:  
-   -------------------------------------------------
 
-Ill-typed arguments supplied to function 'algebra_solver':
-(<F1>, vector, vector, array[] int, array[] int)
-where F1 = (vector, vector, array[] real, array[] int) => vector
-Available signatures:
-(<F2>, vector, vector, data array[] real, data array[] int) => vector
-where F2 = (vector, vector, data array[] real, data array[] int) => vector
-  The fourth argument must be array[] real but got array[] int
-(<F2>, vector, vector, data array[] real, data array[] int, data real,
- data real, data real) => vector
-  Expected 8 arguments but found 5 arguments.
   $ ../../../../../install/default/bin/stanc bad_x_r_type_control.stan
-Semantic error in 'bad_x_r_type_control.stan', line 31, column 10 to column 78:
-   -------------------------------------------------
-    29:  transformed parameters {
-    30:    vector[2] y_s_p;
-    31:    y_s_p = algebra_solver(algebra_system, y, theta_p, x_r, x_i, 0.01, 0.01, 10);
-                   ^
-    32:  }
-    33:  
-   -------------------------------------------------
 
-Ill-typed arguments supplied to function 'algebra_solver':
-(<F1>, vector, vector, array[] int, array[] int, real, real, int)
-where F1 = (vector, vector, array[] real, array[] int) => vector
-Available signatures:
-(<F2>, vector, vector, data array[] real, data array[] int, data real,
- data real, data real) => vector
-where F2 = (vector, vector, data array[] real, data array[] int) => vector
-  The fourth argument must be array[] real but got array[] int
-(<F2>, vector, vector, data array[] real, data array[] int) => vector
-  Expected 5 arguments but found 8 arguments.
   $ ../../../../../install/default/bin/stanc bad_x_r_var_type.stan
 Semantic error in 'bad_x_r_var_type.stan', line 31, column 10 to column 64:
    -------------------------------------------------

test/integration/bad/stanc.expected

@@ -1780,28 +1780,7 @@ Semantic error in 'prob-poisson_log-trunc-low.stan', line 11, column 2 to column
 
 Truncation is only defined if distribution has _lcdf and _lccdf functions implemented with appropriate signature.
   $ ../../../../install/default/bin/stanc propto-bad1.stan
-Semantic error in 'propto-bad1.stan', line 2, column 14 to column 37:
-   -------------------------------------------------
-     1:  model {
-     2:      target += normal_lupdf(1| {0}, 1);
-                       ^
-     3:  }
-   -------------------------------------------------
 
-Ill-typed arguments supplied to function 'normal_lupdf':
-(int, array[] int, int)
-Available signatures:
-(real, real, real) => real
-  The second argument must be real but got array[] int
-(real, real, array[] real) => real
-  The second argument must be real but got array[] int
-(real, real, vector) => real
-  The second argument must be real but got array[] int
-(real, real, row_vector) => real
-  The second argument must be real but got array[] int
-(real, vector, real) => real
-  The second argument must be vector but got array[] int
-(Additional signatures omitted)
   $ ../../../../install/default/bin/stanc real-pdf.stan
 Semantic error in 'real-pdf.stan', line 2, column 2 to line 4, column 3:
    -------------------------------------------------