Call new deserializer backend for constrained reads

src/frontend/Ast_to_Mir.ml

@@ -815,7 +815,9 @@ let trans_block ud_dists declc block prog =
           if Utils.is_user_ident decl_id then
             let constrain_checks =
               match declc.dconstrain with
-              | Some Constrain | Some Unconstrain ->
+              | Some Constrain ->
+                  check_transform_shape decl_id decl_var smeta.loc transform
+              | Some Unconstrain ->
                   check_transform_shape decl_id decl_var smeta.loc transform
                   @ constrain_decl type_ declc.dconstrain transform decl_id
                       decl_var smeta.loc

src/stan_math_backend/Expression_gen.ml

@@ -418,9 +418,20 @@ and pp_compiler_internal_fn ad ut f ppf es =
   | Some FnReadData -> read_data ut ppf es
   | Some FnReadParam -> (
     match es with
-    | {Expr.Fixed.pattern= Lit (Str, base_type); _} :: dims ->
-        pf ppf "@[<hov 2>in__.%s(@,%a)@]" base_type (list ~sep:comma pp_expr)
-          dims
+    | {Expr.Fixed.pattern= Lit (Str, constraint_string); meta= emeta}
+      :: {Expr.Fixed.pattern= Lit (Int, n_constraint_args_str); _} :: args ->
+        let n_constraint_args = int_of_string n_constraint_args_str in
+        let constraint_args, dims = List.split_n args n_constraint_args in
+        let lp_expr = Expr.Fixed.{pattern= Var "lp__"; meta= emeta} in
+        let arg_exprs = constraint_args @ [lp_expr] @ dims in
+        let maybe_constraint, maybe_jacobian =
+          if String.is_empty constraint_string then ("", "")
+          else ("_" ^ constraint_string, ", jacobian__")
+        in
+        pf ppf "@[<hov 2>in__.read%s<%a%s>(@,%a)@]" maybe_constraint
+          pp_unsizedtype_local
+          (UnsizedType.AutoDiffable, ut)
+          maybe_jacobian (list ~sep:comma pp_expr) arg_exprs
     | _ -> raise_s [%message "emit ReadParam with " (es : Expr.Typed.t list)] )
   | _ -> gen_fun_app ppf f es
 

src/stan_math_backend/Stan_math_code_gen.ml

@@ -510,7 +510,7 @@ let pp_write_array ppf {Program.prog_name; generate_quantities; _} =
   let intro ppf () =
     pf ppf "%a@ %a@ %a" (list ~sep:cut string)
       [ "using local_scalar_t__ = double;"; "vars__.resize(0);"
-      ; "stan::io::reader<local_scalar_t__> in__(params_r__, params_i__);"
+      ; "stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);"
       ; "double lp__ = 0.0;"
       ; "(void) lp__;  // dummy to suppress unused var warning"
       ; "int current_statement__ = 0; "

src/stan_math_backend/Transform_Mir.ml

@@ -161,56 +161,80 @@ let data_read smeta (decl_id, st) =
           ; Stmt.Helpers.for_scalar_inv st bodyfn decl_var smeta ]
         |> swrap ]
 
-let rec base_ut_to_string = function
-  | UnsizedType.UMatrix -> "matrix"
-  | UVector -> "vector"
-  | URowVector -> "row_vector"
-  | UReal -> "scalar"
-  | UInt -> "integer"
-  | UArray t -> base_ut_to_string t
-  | t ->
-      raise_s
-        [%message "Another place where it's weird to get " (t : UnsizedType.t)]
+type constrainaction = Check | Constrain | Unconstrain [@@deriving sexp]
+
+let check_constraint_to_string t (c : constrainaction) =
+  match t with
+  | Program.Ordered -> "ordered"
+  | PositiveOrdered -> "positive_ordered"
+  | Simplex -> "simplex"
+  | UnitVector -> "unit_vector"
+  | CholeskyCorr -> "cholesky_factor_corr"
+  | CholeskyCov -> "cholesky_factor"
+  | Correlation -> "corr_matrix"
+  | Covariance -> "cov_matrix"
+  | Lower _ -> (
+    match c with
+    | Check -> "greater_or_equal"
+    | Constrain | Unconstrain -> "lb" )
+  | Upper _ -> (
+    match c with Check -> "less_or_equal" | Constrain | Unconstrain -> "ub" )
+  | LowerUpper _ -> (
+    match c with
+    | Check ->
+        raise_s
+          [%message "LowerUpper is really two other checks tied together"]
+    | Constrain | Unconstrain -> "lub" )
+  | Offset _ | Multiplier _ | OffsetMultiplier _ -> (
+    match c with Check -> "" | Constrain | Unconstrain -> "offset_multiplier" )
+  | Identity -> ""
+
+let constrain_constraint_to_string t (c : constrainaction) =
+  match t with
+  | Program.CholeskyCorr -> "cholesky_corr"
+  | _ -> check_constraint_to_string t c
+
+let default_multiplier = 1
+let default_offset = 0
+
+let transform_args = function
+  | Program.Offset offset -> [offset; Expr.Helpers.int default_multiplier]
+  | Multiplier multiplier -> [Expr.Helpers.int default_offset; multiplier]
+  | transform ->
+      Program.fold_transformation (fun args arg -> args @ [arg]) [] transform
 
 let param_read smeta
-    ( decl_id
-    , Program.({ out_constrained_st= cst
-               ; out_unconstrained_st= ucst
-               ; out_block; _ }) ) =
+    (decl_id, Program.({out_constrained_st= cst; out_block; out_trans; _})) =
   if not (out_block = Parameters) then []
   else
-    let decl_id, decl =
-      match cst = ucst with
-      | true -> (decl_id, [])
-      | false ->
-          let decl_id = decl_id ^ "_in__" in
-          let d =
-            Stmt.Fixed.Pattern.Decl
-              {decl_adtype= AutoDiffable; decl_id; decl_type= Sized ucst}
-          in
-          (decl_id, [Stmt.Fixed.{meta= smeta; pattern= d}])
+    let ut = SizedType.to_unsized cst in
+    let decl_var =
+      Expr.Fixed.
+        { pattern= Var decl_id
+        ; meta=
+            Expr.Typed.Meta.create ~loc:smeta ~type_:ut ~adlevel:AutoDiffable
+              () }
     in
-    let unconstrained_decl_var =
-      let meta =
-        Expr.Typed.Meta.create ~loc:smeta
-          ~type_:SizedType.(to_unsized cst)
-          ~adlevel:AutoDiffable ()
-      in
-      Expr.Fixed.{meta; pattern= Var decl_id}
-    in
-    let bodyfn var =
-      let readfnapp (var : Expr.Typed.t) =
-        Expr.(
-          Helpers.(
-            internal_funapp FnReadParam
-              ( str (base_ut_to_string (SizedType.to_unsized ucst))
-              :: SizedType.dims_of ucst ))
-            Typed.Meta.{var.meta with type_= base_type ucst})
-      in
-      Stmt.Helpers.assign_indexed (SizedType.to_unsized cst) decl_id smeta
-        readfnapp var
+    let transform_args = transform_args out_trans in
+    (* this is an absolute hack
+
+       I need to unpack the constraint arguments and the dimensions in codegen, but we pack them all together into a fake function FnReadParam as expressions
+       So, I'm packing in the number of constraint arguments to read as an int expression
+       To avoid this hack, keep internal functions as a variant type instead of a normal funapp
+    *)
+    let n_args_expression = Expr.Helpers.int (List.length transform_args) in
+    let read =
+      Expr.(
+        Helpers.(
+          internal_funapp FnReadParam
+            ( Expr.Helpers.str
+                (constrain_constraint_to_string out_trans Constrain)
+            :: n_args_expression
+            :: (transform_args @ SizedType.get_dims cst) ))
+          Typed.Meta.{decl_var.meta with type_= ut})
     in
-    decl @ [Stmt.Helpers.for_eigen ucst bodyfn unconstrained_decl_var smeta]
+    [ Stmt.Fixed.
+        {pattern= Pattern.Assignment ((decl_id, ut, []), read); meta= smeta} ]
 
 let escape_name str =
   str

test/integration/cli-args/filename_good.expected

@@ -121,7 +121,7 @@ class filename_good_model final : public model_base_crtp<filename_good_model> {
                                std::ostream* pstream__ = nullptr) const {
     using local_scalar_t__ = double;
     vars__.resize(0);
-    stan::io::reader<local_scalar_t__> in__(params_r__, params_i__);
+    stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);
     double lp__ = 0.0;
     (void) lp__;  // dummy to suppress unused var warning
     int current_statement__ = 0; 

test/integration/good/code-gen/cl.expected

@@ -579,60 +579,64 @@ class optimize_glm_model final : public model_base_crtp<optimize_glm_model> {
       stan::math::fill(alpha_v, DUMMY_VAR__);
 
       current_statement__ = 1;
-      alpha_v = in__.vector(k);
+      assign(alpha_v, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable alpha_v");
       Eigen::Matrix<local_scalar_t__, -1, 1> beta;
       beta = Eigen::Matrix<local_scalar_t__, -1, 1>(k);
       stan::math::fill(beta, DUMMY_VAR__);
 
       current_statement__ = 2;
-      beta = in__.vector(k);
+      assign(beta, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable beta");
       Eigen::Matrix<local_scalar_t__, -1, 1> cuts;
       cuts = Eigen::Matrix<local_scalar_t__, -1, 1>(k);
       stan::math::fill(cuts, DUMMY_VAR__);
 
       current_statement__ = 3;
-      cuts = in__.vector(k);
+      assign(cuts, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable cuts");
       local_scalar_t__ sigma;
       sigma = DUMMY_VAR__;
 
       current_statement__ = 4;
-      sigma = in__.scalar();
-      current_statement__ = 4;
-      if (jacobian__) {
-        current_statement__ = 4;
-        sigma = stan::math::lb_constrain(sigma, 0, lp__);
-      } else {
-        current_statement__ = 4;
-        sigma = stan::math::lb_constrain(sigma, 0);
-      }
+      sigma = in__.read_lb<local_scalar_t__, jacobian__>(0, lp__);
       local_scalar_t__ alpha;
       alpha = DUMMY_VAR__;
 
       current_statement__ = 5;
-      alpha = in__.scalar();
+      alpha = in__.read<local_scalar_t__>(lp__);
       local_scalar_t__ phi;
       phi = DUMMY_VAR__;
 
       current_statement__ = 6;
-      phi = in__.scalar();
+      phi = in__.read<local_scalar_t__>(lp__);
       Eigen::Matrix<local_scalar_t__, -1, -1> X_p;
       X_p = Eigen::Matrix<local_scalar_t__, -1, -1>(n, k);
       stan::math::fill(X_p, DUMMY_VAR__);
 
       current_statement__ = 7;
-      X_p = in__.matrix(n, k);
+      assign(X_p, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, -1>>(lp__, n, k),
+        "assigning variable X_p");
       Eigen::Matrix<local_scalar_t__, -1, -1> beta_m;
       beta_m = Eigen::Matrix<local_scalar_t__, -1, -1>(n, k);
       stan::math::fill(beta_m, DUMMY_VAR__);
 
       current_statement__ = 8;
-      beta_m = in__.matrix(n, k);
+      assign(beta_m, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, -1>>(lp__, n, k),
+        "assigning variable beta_m");
       Eigen::Matrix<local_scalar_t__, 1, -1> X_rv_p;
       X_rv_p = Eigen::Matrix<local_scalar_t__, 1, -1>(n);
       stan::math::fill(X_rv_p, DUMMY_VAR__);
 
       current_statement__ = 9;
-      X_rv_p = in__.row_vector(n);
+      assign(X_rv_p, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, 1, -1>>(lp__, n),
+        "assigning variable X_rv_p");
       {
         current_statement__ = 10;
         lp_accum__.add(
@@ -1288,7 +1292,7 @@ class optimize_glm_model final : public model_base_crtp<optimize_glm_model> {
                                std::ostream* pstream__ = nullptr) const {
     using local_scalar_t__ = double;
     vars__.resize(0);
-    stan::io::reader<local_scalar_t__> in__(params_r__, params_i__);
+    stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);
     double lp__ = 0.0;
     (void) lp__;  // dummy to suppress unused var warning
     int current_statement__ = 0; 
@@ -1304,54 +1308,64 @@ class optimize_glm_model final : public model_base_crtp<optimize_glm_model> {
       stan::math::fill(alpha_v, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 1;
-      alpha_v = in__.vector(k);
+      assign(alpha_v, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable alpha_v");
       Eigen::Matrix<double, -1, 1> beta;
       beta = Eigen::Matrix<double, -1, 1>(k);
       stan::math::fill(beta, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 2;
-      beta = in__.vector(k);
+      assign(beta, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable beta");
       Eigen::Matrix<double, -1, 1> cuts;
       cuts = Eigen::Matrix<double, -1, 1>(k);
       stan::math::fill(cuts, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 3;
-      cuts = in__.vector(k);
+      assign(cuts, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, 1>>(lp__, k),
+        "assigning variable cuts");
       double sigma;
       sigma = std::numeric_limits<double>::quiet_NaN();
 
       current_statement__ = 4;
-      sigma = in__.scalar();
-      current_statement__ = 4;
-      sigma = stan::math::lb_constrain(sigma, 0);
+      sigma = in__.read_lb<local_scalar_t__, jacobian__>(0, lp__);
       double alpha;
       alpha = std::numeric_limits<double>::quiet_NaN();
 
       current_statement__ = 5;
-      alpha = in__.scalar();
+      alpha = in__.read<local_scalar_t__>(lp__);
       double phi;
       phi = std::numeric_limits<double>::quiet_NaN();
 
       current_statement__ = 6;
-      phi = in__.scalar();
+      phi = in__.read<local_scalar_t__>(lp__);
       Eigen::Matrix<double, -1, -1> X_p;
       X_p = Eigen::Matrix<double, -1, -1>(n, k);
       stan::math::fill(X_p, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 7;
-      X_p = in__.matrix(n, k);
+      assign(X_p, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, -1>>(lp__, n, k),
+        "assigning variable X_p");
       Eigen::Matrix<double, -1, -1> beta_m;
       beta_m = Eigen::Matrix<double, -1, -1>(n, k);
       stan::math::fill(beta_m, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 8;
-      beta_m = in__.matrix(n, k);
+      assign(beta_m, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, -1, -1>>(lp__, n, k),
+        "assigning variable beta_m");
       Eigen::Matrix<double, 1, -1> X_rv_p;
       X_rv_p = Eigen::Matrix<double, 1, -1>(n);
       stan::math::fill(X_rv_p, std::numeric_limits<double>::quiet_NaN());
 
       current_statement__ = 9;
-      X_rv_p = in__.row_vector(n);
+      assign(X_rv_p, nil_index_list(),
+        in__.read<Eigen::Matrix<local_scalar_t__, 1, -1>>(lp__, n),
+        "assigning variable X_rv_p");
       for (int sym1__ = 1; sym1__ <= k; ++sym1__) {
         vars__.emplace_back(alpha_v[(sym1__ - 1)]);}
       for (int sym1__ = 1; sym1__ <= k; ++sym1__) {