diff --git a/.ocamlformat b/.ocamlformat index 6bafbd1c4e..3b217634ab 100644 --- a/.ocamlformat +++ b/.ocamlformat @@ -1 +1 @@ -profile=compact +profile=janestreet diff --git a/scripts/install_build_deps.sh b/scripts/install_build_deps.sh index f22f46d1dd..6cc2932f34 100644 --- a/scripts/install_build_deps.sh +++ b/scripts/install_build_deps.sh @@ -1,6 +1,6 @@ #!/bin/bash # Menhir is our parsing library and annoyingly its module name does not match # its library name, so we install it manually here. -opam pin -y ocamlformat 0.8 +opam pin -y ocamlformat 0.13 opam pin -y dune 1.11.4 opam install -y core_kernel.v0.11.1 menhir.20181113 ppx_deriving.4.2.1 fmt.0.8.5 yojson.1.7.0 diff --git a/scripts/install_dev_deps.sh b/scripts/install_dev_deps.sh index 3324563b7d..0c63aaf74d 100644 --- a/scripts/install_dev_deps.sh +++ b/scripts/install_dev_deps.sh @@ -1,3 +1,3 @@ #!/bin/bash # Merlin, utop, ocp-indent, ocamlformat, and patdiff are all for developer assistance -opam install -y ocamlformat.0.8 merlin utop ocp-indent patdiff +opam install -y ocamlformat.0.13 merlin utop ocp-indent patdiff diff --git a/src/analysis_and_optimization/Dataflow_types.ml b/src/analysis_and_optimization/Dataflow_types.ml index 9a80251dfd..9bb6aea2f0 100644 --- a/src/analysis_and_optimization/Dataflow_types.ml +++ b/src/analysis_and_optimization/Dataflow_types.ml @@ -30,7 +30,10 @@ type reaching_defn = vexpr * label [@@deriving sexp, hash, compare] type source_loc = | MirNode of Location_span.t | StartOfBlock - | TargetTerm of {term: Expr.Typed.t; assignment_label: label} + | TargetTerm of + { term : Expr.Typed.t + ; assignment_label : label + } [@@deriving sexp] (** @@ -49,11 +52,12 @@ type source_loc = there is none *) type 'rd_info node_info = - { rd_sets: 'rd_info - ; possible_previous: label Set.Poly.t - ; rhs_set: vexpr Set.Poly.t - ; controlflow: label Set.Poly.t - ; loc: source_loc } + { rd_sets : 'rd_info + ; possible_previous : label Set.Poly.t + ; rhs_set : vexpr Set.Poly.t + ; controlflow : label Set.Poly.t + ; loc : source_loc + } [@@deriving sexp] (** @@ -61,8 +65,7 @@ type 'rd_info node_info = function that maps from the 'entry' set to the 'exit' set, where the entry set is what's true before executing this node and the exit set is true after. *) -type node_info_update = - (reaching_defn Set.Poly.t -> reaching_defn Set.Poly.t) node_info +type node_info_update = (reaching_defn Set.Poly.t -> reaching_defn Set.Poly.t) node_info (** A node_info where the reaching definition information is explicitly written as the @@ -84,14 +87,15 @@ type node_info_fixedpoint = * returns: A set of the return nodes that have been encountered *) type traversal_state = - { label_ix: label - ; node_info_map: (int, node_info_update) Map.Poly.t - ; possible_previous: label Set.Poly.t - ; target_terms: label Set.Poly.t - ; continues: label Set.Poly.t - ; breaks: label Set.Poly.t - ; returns: label Set.Poly.t - ; rejects: label Set.Poly.t } + { label_ix : label + ; node_info_map : (int, node_info_update) Map.Poly.t + ; possible_previous : label Set.Poly.t + ; target_terms : label Set.Poly.t + ; continues : label Set.Poly.t + ; breaks : label Set.Poly.t + ; returns : label Set.Poly.t + ; rejects : label Set.Poly.t + } (** The most recently nested control flow (block start, if/then, or loop) @@ -109,9 +113,10 @@ type cf_state = label excluded for non-statistical dependency analysis *) type dataflow_graph = - { node_info_map: (int, node_info_fixedpoint) Map.Poly.t - ; possible_exits: label Set.Poly.t - ; probabilistic_nodes: label Set.Poly.t } + { node_info_map : (int, node_info_fixedpoint) Map.Poly.t + ; possible_exits : label Set.Poly.t + ; probabilistic_nodes : label Set.Poly.t + } [@@deriving sexp] (** @@ -120,5 +125,8 @@ type dataflow_graph = See Middle.prog for block descriptions. *) type prog_df_graphs = - {tdatab: dataflow_graph; modelb: dataflow_graph; gqb: dataflow_graph} + { tdatab : dataflow_graph + ; modelb : dataflow_graph + ; gqb : dataflow_graph + } [@@deriving sexp] diff --git a/src/analysis_and_optimization/Dataflow_utils.ml b/src/analysis_and_optimization/Dataflow_utils.ml index 3f4f268362..e050c70404 100644 --- a/src/analysis_and_optimization/Dataflow_utils.ml +++ b/src/analysis_and_optimization/Dataflow_utils.ml @@ -4,8 +4,9 @@ open Dataflow_types open Mir_utils (** Union maps, preserving the left element in a collision *) -let union_maps_left (m1 : ('a, 'b) Map.Poly.t) (m2 : ('a, 'b) Map.Poly.t) : - ('a, 'b) Map.Poly.t = +let union_maps_left (m1 : ('a, 'b) Map.Poly.t) (m2 : ('a, 'b) Map.Poly.t) + : ('a, 'b) Map.Poly.t + = let f ~key:_ opt = match opt with | `Left v -> Some v @@ -13,6 +14,7 @@ let union_maps_left (m1 : ('a, 'b) Map.Poly.t) (m2 : ('a, 'b) Map.Poly.t) : | `Both (v1, _) -> Some v1 in Map.Poly.merge m1 m2 ~f +;; (** Merge two maps whose values are sets, and union the sets when there's a collision. @@ -25,13 +27,15 @@ let merge_set_maps m1 m2 = | `Both (e1, e2) -> Some (Set.Poly.union e1 e2) in Map.Poly.merge ~f:merge_map_elems m1 m2 +;; (** Generate a Map by applying a function to each element of a key set. *) let generate_map s ~f = Set.Poly.fold s ~init:Map.Poly.empty ~f:(fun m e -> - Map.Poly.add_exn m ~key:e ~data:(f e) ) + Map.Poly.add_exn m ~key:e ~data:(f e)) +;; (** Like a forward traversal, but branches accumulate two different states that are @@ -41,19 +45,21 @@ let branching_traverse_statement stmt ~join ~init ~f = Stmt.Fixed.Pattern.( match stmt with | IfElse (pred, then_s, else_s_opt) -> - let s', c = f init then_s in - Option.value_map else_s_opt - ~default:(join s' init, IfElse (pred, c, None)) - ~f:(fun else_s -> - let s'', c' = f init else_s in - (join s' s'', IfElse (pred, c, Some c')) ) + let s', c = f init then_s in + Option.value_map + else_s_opt + ~default:(join s' init, IfElse (pred, c, None)) + ~f:(fun else_s -> + let s'', c' = f init else_s in + join s' s'', IfElse (pred, c, Some c')) | _ as s -> fwd_traverse_statement s ~init ~f) +;; (** Like a branching traversal, but doesn't return an updated statement. *) let branching_fold_statement stmt ~join ~init ~f = - fst - (branching_traverse_statement stmt ~join ~init ~f:(fun s a -> (f s a, ()))) + fst (branching_traverse_statement stmt ~join ~init ~f:(fun s a -> f s a, ())) +;; (** See interface file @@ -67,12 +73,12 @@ let build_statement_map extract metadata stmt = fwd_traverse_statement (extract stmt) ~init:(next_label', map) ~f in ( ( next_label'' - , union_maps_left map - (Map.Poly.singleton this_label (built, metadata stmt)) ) + , union_maps_left map (Map.Poly.singleton this_label (built, metadata stmt)) ) , this_label ) in let (_, map), _ = build_statement_map_rec 1 Map.Poly.empty stmt in map +;; (* TODO: this currently does not seem to be labelling inside function bodies. Could we also do that? *) @@ -85,6 +91,7 @@ let rec build_recursive_statement rebuild statement_map label = let build_stmt = build_recursive_statement rebuild statement_map in let stmt = Stmt.Fixed.Pattern.map Fn.id build_stmt stmt_ints in rebuild stmt meta +;; (** Represents the state required to build control flow information during an MIR traversal, where @@ -95,25 +102,31 @@ let rec build_recursive_statement rebuild statement_map label = node *) type cf_state = - { breaks: label Set.Poly.t - ; continues: label Set.Poly.t - ; returns: label Set.Poly.t - ; exits: label Set.Poly.t } + { breaks : label Set.Poly.t + ; continues : label Set.Poly.t + ; returns : label Set.Poly.t + ; exits : label Set.Poly.t + } (** Represents the control flow information at each node in the control graph, where * predecessors points to the nodes which could have executed before this node * parents points to the adjacent nodes which directly influence the execution of this node *) -type cf_edges = {predecessors: label Set.Poly.t; parents: label Set.Poly.t} +type cf_edges = + { predecessors : label Set.Poly.t + ; parents : label Set.Poly.t + } (** Join the state of a controlflow traversal across different branches of execution such as over if/else branch. *) let join_cf_states (state1 : cf_state) (state2 : cf_state) : cf_state = - { breaks= Set.Poly.union state1.breaks state2.breaks - ; continues= Set.Poly.union state1.continues state2.continues - ; returns= Set.Poly.union state1.returns state2.returns - ; exits= Set.Poly.union state1.exits state2.exits } + { breaks = Set.Poly.union state1.breaks state2.breaks + ; continues = Set.Poly.union state1.continues state2.continues + ; returns = Set.Poly.union state1.returns state2.returns + ; exits = Set.Poly.union state1.exits state2.exits + } +;; (** Check if the statement controls the execution of its substatements. *) let is_ctrl_flow pattern = @@ -122,22 +135,25 @@ let is_ctrl_flow pattern = | While _ -> true | For _ -> true | _ -> false +;; (** Simultaneously builds the controlflow parent graph, the predecessor graph and the exit set of a statement. It's advantageous to build them together because they both rely on some of the same Break, Continue and Return bookkeeping. *) -let build_cf_graphs ?(flatten_loops = false) ?(blocks_after_body = true) - statement_map = - let rec build_cf_graph_rec (cf_parent : label option) +let build_cf_graphs ?(flatten_loops = false) ?(blocks_after_body = true) statement_map = + let rec build_cf_graph_rec + (cf_parent : label option) ((in_state, in_map) : cf_state * (label, cf_edges) Map.Poly.t) - (label : label) : cf_state * (label, cf_edges) Map.Poly.t = + (label : label) + : cf_state * (label, cf_edges) Map.Poly.t + = let stmt, _ = Map.Poly.find_exn statement_map label in (* Only control flow nodes should pass themselves down as parents *) let child_cf = if is_ctrl_flow stmt then Some label else cf_parent in let join (state1, map1) (state2, map2) = - (join_cf_states state1 state2, union_maps_left map1 map2) + join_cf_states state1 state2, union_maps_left map1 map2 in (* This node is the parent of substatements, unless this is a Block, which is visited after substatements *) @@ -148,8 +164,10 @@ let build_cf_graphs ?(flatten_loops = false) ?(blocks_after_body = true) in (* The accumulated state after traversing substatements *) let substmt_state_unlooped, substmt_map = - branching_fold_statement stmt ~join - ~init:({in_state with exits= substmt_preds}, in_map) + branching_fold_statement + stmt + ~join + ~init:({ in_state with exits = substmt_preds }, in_map) ~f:(build_cf_graph_rec child_cf) in (* If the statement is a loop, we need to include the loop body exits as predecessors @@ -157,43 +175,48 @@ let build_cf_graphs ?(flatten_loops = false) ?(blocks_after_body = true) let substmt_state, predecessors = match stmt with | For _ | While _ -> - (* Loop statements are preceded by: + (* Loop statements are preceded by: 1. The statements that come before the loop 2. The natural exit points of the loop body 3. Continue statements in the loop body This comment mangling brought to you by the autoformatter *) - let loop_predecessors = - Set.Poly.union_list - [ (*1*) in_state.exits; (*2*) substmt_state_unlooped.exits - ; (*3*) - Set.Poly.diff substmt_state_unlooped.continues - in_state.continues ] - in - (* Loop exits are: + let loop_predecessors = + Set.Poly.union_list + [ (*1*) + in_state.exits + ; (*2*) + substmt_state_unlooped.exits + ; (*3*) + Set.Poly.diff substmt_state_unlooped.continues in_state.continues + ] + in + (* Loop exits are: 1. The loop node itself, since the last action of a typical loop execution is to check if there are any iterations remaining 2. Break statements in the loop body, since broken loops don't execute the loop statement *) - let loop_exits = - if flatten_loops then substmt_state_unlooped.exits - else - Set.Poly.union_list - [ (*1*) Set.Poly.singleton label - ; (*2*) - Set.Poly.diff substmt_state_unlooped.breaks in_state.breaks - ] - in - ({substmt_state_unlooped with exits= loop_exits}, loop_predecessors) + let loop_exits = + if flatten_loops + then substmt_state_unlooped.exits + else + Set.Poly.union_list + [ (*1*) + Set.Poly.singleton label + ; (*2*) + Set.Poly.diff substmt_state_unlooped.breaks in_state.breaks + ] + in + { substmt_state_unlooped with exits = loop_exits }, loop_predecessors | Block _ when blocks_after_body -> - (* Block statements are preceded by the natural exit points of the block + (* Block statements are preceded by the natural exit points of the block body *) - let block_predecessors = substmt_state_unlooped.exits in - (* Block exits are just the block node *) - let block_exits = Set.Poly.singleton label in - ({substmt_state_unlooped with exits= block_exits}, block_predecessors) - | _ -> (substmt_state_unlooped, in_state.exits) + let block_predecessors = substmt_state_unlooped.exits in + (* Block exits are just the block node *) + let block_exits = Set.Poly.singleton label in + { substmt_state_unlooped with exits = block_exits }, block_predecessors + | _ -> substmt_state_unlooped, in_state.exits in (* Some statements interact with the break/return/continue states E.g., loops nullify breaks and continues in their body, but are still affected by @@ -201,66 +224,77 @@ let build_cf_graphs ?(flatten_loops = false) ?(blocks_after_body = true) let breaks_out, returns_out, continues_out, extra_cf_deps = match stmt with | Break -> - ( Set.Poly.add substmt_state.breaks label - , substmt_state.returns - , substmt_state.continues - , Set.Poly.empty ) + ( Set.Poly.add substmt_state.breaks label + , substmt_state.returns + , substmt_state.continues + , Set.Poly.empty ) | Return _ -> - ( substmt_state.breaks - , Set.Poly.add substmt_state.returns label - , substmt_state.continues - , Set.Poly.empty ) + ( substmt_state.breaks + , Set.Poly.add substmt_state.returns label + , substmt_state.continues + , Set.Poly.empty ) | Continue -> - ( substmt_state.breaks - , substmt_state.returns - , Set.Poly.add substmt_state.continues label - , Set.Poly.empty ) + ( substmt_state.breaks + , substmt_state.returns + , Set.Poly.add substmt_state.continues label + , Set.Poly.empty ) | While _ | For _ -> - ( in_state.breaks - , substmt_state.returns - , in_state.continues - , Set.Poly.union substmt_state.breaks substmt_state.returns ) + ( in_state.breaks + , substmt_state.returns + , in_state.continues + , Set.Poly.union substmt_state.breaks substmt_state.returns ) | _ -> - ( substmt_state.breaks - , substmt_state.returns - , substmt_state.continues - , Set.Poly.empty ) + ( substmt_state.breaks + , substmt_state.returns + , substmt_state.continues + , Set.Poly.empty ) in let cf_parents = Set.Poly.union_list - [ Option.value_map cf_parent ~default:Set.Poly.empty - ~f:Set.Poly.singleton - ; in_state.returns; in_state.continues; extra_cf_deps ] + [ Option.value_map cf_parent ~default:Set.Poly.empty ~f:Set.Poly.singleton + ; in_state.returns + ; in_state.continues + ; extra_cf_deps + ] in - ( { breaks= breaks_out - ; continues= continues_out - ; returns= returns_out - ; exits= substmt_state.exits } - , Map.Poly.add_exn substmt_map ~key:label - ~data:{parents= cf_parents; predecessors} ) + ( { breaks = breaks_out + ; continues = continues_out + ; returns = returns_out + ; exits = substmt_state.exits + } + , Map.Poly.add_exn substmt_map ~key:label ~data:{ parents = cf_parents; predecessors } + ) in let state, edges = - build_cf_graph_rec None - ( { breaks= Set.Poly.empty - ; continues= Set.Poly.empty - ; returns= Set.Poly.empty - ; exits= Set.Poly.empty } + build_cf_graph_rec + None + ( { breaks = Set.Poly.empty + ; continues = Set.Poly.empty + ; returns = Set.Poly.empty + ; exits = Set.Poly.empty + } , Map.Poly.empty ) 1 in ( state.exits , Map.Poly.map edges ~f:(fun e -> e.predecessors) , Map.Poly.map edges ~f:(fun e -> e.parents) ) +;; (** See interface file *) let build_cf_graph statement_map = let _, _, cf_graph = build_cf_graphs statement_map in cf_graph +;; (** See interface file *) -let build_predecessor_graph ?(flatten_loops = false) - ?(blocks_after_body = true) statement_map = +let build_predecessor_graph + ?(flatten_loops = false) + ?(blocks_after_body = true) + statement_map + = let exits, pred_graph, _ = build_cf_graphs ~flatten_loops ~blocks_after_body statement_map in - (exits, pred_graph) + exits, pred_graph +;; diff --git a/src/analysis_and_optimization/Dataflow_utils.mli b/src/analysis_and_optimization/Dataflow_utils.mli index 0d23f5e28a..39143fd748 100644 --- a/src/analysis_and_optimization/Dataflow_utils.mli +++ b/src/analysis_and_optimization/Dataflow_utils.mli @@ -2,17 +2,9 @@ open Core_kernel open Middle open Dataflow_types -val union_maps_left : - ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t (** Union maps, preserving the left element in a collision *) +val union_maps_left : ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t -val build_cf_graphs : - ?flatten_loops:bool - -> ?blocks_after_body:bool - -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t - -> label Set.Poly.t - * (label, label Set.Poly.t) Map.Poly.t - * (label, label Set.Poly.t) Map.Poly.t (** Simultaneously builds the controlflow parent graph, the predecessor graph and the exit set of a statement. It's advantageous to build them together because they both rely on @@ -24,22 +16,24 @@ val build_cf_graphs : * (exit set, predecessor graph) is the return value of build_predecessor_graph * (controlflow parent graph) is the return value of build_cf_graph *) +val build_cf_graphs + : ?flatten_loops:bool + -> ?blocks_after_body:bool + -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t + -> label Set.Poly.t + * (label, label Set.Poly.t) Map.Poly.t + * (label, label Set.Poly.t) Map.Poly.t -val build_cf_graph : - (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t - -> (label, label Set.Poly.t) Map.Poly.t (** Building the controlflow graph requires a traversal with state that includes continues, breaks, returns and the controlflow graph accumulator. The traversal should be a branching traversal with set unions rather than a forward traversal because continue and return statements shouldn't affect other branches of execution. *) +val build_cf_graph + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t + -> (label, label Set.Poly.t) Map.Poly.t -val build_predecessor_graph : - ?flatten_loops:bool - -> ?blocks_after_body:bool - -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t - -> label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t (** Building the predecessor graph requires a traversal with state that includes the current previous nodes and the predecessor graph accumulator. Special cases are made @@ -47,41 +41,46 @@ val build_predecessor_graph : they should include loop predecessors in their exit sets. I'm not sure if the single re-traversal of the loop body is sufficient or this requires finding a fixed-point. *) +val build_predecessor_graph + : ?flatten_loops:bool + -> ?blocks_after_body:bool + -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t + -> label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t -val build_recursive_statement : - (('e, 's) Stmt.Fixed.Pattern.t -> 'm -> 's) - -> (label, ('e, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t - -> label - -> 's (** Build a fixed-point data type representation of a statement given a label-map representation. *) +val build_recursive_statement + : (('e, 's) Stmt.Fixed.Pattern.t -> 'm -> 's) + -> (label, ('e, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t + -> label + -> 's -val is_ctrl_flow : ('a, 'b) Stmt.Fixed.Pattern.t -> bool (** Check if the statement controls the execution of its substatements. *) +val is_ctrl_flow : ('a, 'b) Stmt.Fixed.Pattern.t -> bool -val merge_set_maps : - ('a, 'b Set.Poly.t) Map.Poly.t - -> ('a, 'b Set.Poly.t) Map.Poly.t - -> ('a, 'b Set.Poly.t) Map.Poly.t (** Merge two maps whose values are sets, and union the sets when there's a collision. *) +val merge_set_maps + : ('a, 'b Set.Poly.t) Map.Poly.t + -> ('a, 'b Set.Poly.t) Map.Poly.t + -> ('a, 'b Set.Poly.t) Map.Poly.t -val generate_map : 'a Set.Poly.t -> f:('a -> 'b) -> ('a, 'b) Map.Poly.t (** Generate a Map by applying a function to each element of a key set. *) +val generate_map : 'a Set.Poly.t -> f:('a -> 'b) -> ('a, 'b) Map.Poly.t -val build_statement_map : - ('s -> ('e, 's) Stmt.Fixed.Pattern.t) - -> ('s -> 'm) - -> 's - -> (label, ('e, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t (** The statement map is built by traversing substatements recursively to replace substatements with their labels while building up the substatements' statement maps. Then, the result is the union of the substatement maps with this statement's singleton pair, which is expressed in terms of the new label-containing statement. *) +val build_statement_map + : ('s -> ('e, 's) Stmt.Fixed.Pattern.t) + -> ('s -> 'm) + -> 's + -> (label, ('e, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t diff --git a/src/analysis_and_optimization/Dependence_analysis.ml b/src/analysis_and_optimization/Dependence_analysis.ml index d76e77cc4a..59f3b6eef1 100644 --- a/src/analysis_and_optimization/Dependence_analysis.ml +++ b/src/analysis_and_optimization/Dependence_analysis.ml @@ -11,25 +11,28 @@ open Monotone_framework (***********************************) type node_dep_info = - { predecessors: label Set.Poly.t - ; parents: label Set.Poly.t - ; reaching_defn_entry: reaching_defn Set.Poly.t - ; reaching_defn_exit: reaching_defn Set.Poly.t - ; meta: Location_span.t } + { predecessors : label Set.Poly.t + ; parents : label Set.Poly.t + ; reaching_defn_entry : reaching_defn Set.Poly.t + ; reaching_defn_exit : reaching_defn Set.Poly.t + ; meta : Location_span.t + } (** Find all of the reaching definitions of a variable in an RD set *) -let reaching_defn_lookup (rds : reaching_defn Set.Poly.t) (var : vexpr) : - label Set.Poly.t = +let reaching_defn_lookup (rds : reaching_defn Set.Poly.t) (var : vexpr) : label Set.Poly.t + = Set.Poly.map (Set.Poly.filter rds ~f:(fun (var', _) -> var' = var)) ~f:snd +;; let node_immediate_dependencies (statement_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) - (label : label) : label Set.Poly.t = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) + (label : label) + : label Set.Poly.t + = let stmt, info = Map.Poly.find_exn statement_map label in let rhs_set = Set.Poly.map (stmt_rhs_var_set stmt) ~f:fst in let rhs_deps = @@ -38,6 +41,7 @@ let node_immediate_dependencies ~f:(reaching_defn_lookup info.reaching_defn_entry) in Set.Poly.union info.parents rhs_deps +;; (* This is doing an explicit graph traversal with edges defined by @@ -45,29 +49,37 @@ let node_immediate_dependencies *) let rec node_dependencies_rec (statement_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) - (visited : label Set.Poly.t) (label : label) : label Set.Poly.t = - if Set.Poly.mem visited label then visited - else + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) + (visited : label Set.Poly.t) + (label : label) + : label Set.Poly.t + = + if Set.Poly.mem visited label + then visited + else ( let visited' = Set.Poly.add visited label in let deps = node_immediate_dependencies statement_map ~blockers label in - Set.Poly.fold deps ~init:visited' ~f:(node_dependencies_rec statement_map) + Set.Poly.fold deps ~init:visited' ~f:(node_dependencies_rec statement_map)) +;; let node_dependencies (statement_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) (label : label) : label Set.Poly.t = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + (label : label) + : label Set.Poly.t + = node_dependencies_rec statement_map Set.Poly.empty label +;; let node_vars_dependencies (statement_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) - (vars : vexpr Set.Poly.t) (label : label) : label Set.Poly.t = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + ?(blockers : vexpr Set.Poly.t = Set.Poly.empty) + (vars : vexpr Set.Poly.t) + (label : label) + : label Set.Poly.t + = let _, info = Map.Poly.find_exn statement_map label in let var_deps = union_map @@ -78,6 +90,7 @@ let node_vars_dependencies (Set.union info.parents var_deps) ~init:Set.Poly.empty ~f:(node_dependencies_rec statement_map ~blockers) +;; (* The strategy here is to write an update function on the whole dependency graph in terms @@ -87,53 +100,49 @@ let node_vars_dependencies *) let all_node_dependencies (statement_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) : (label, label Set.Poly.t) Map.Poly.t = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + : (label, label Set.Poly.t) Map.Poly.t + = let immediate_map = Map.mapi statement_map ~f:(fun ~key:label ~data:_ -> - node_immediate_dependencies statement_map label ) + node_immediate_dependencies statement_map label) in let step_node label m = let immediate = Map.find_exn immediate_map label in let updated = - Set.union - (union_map immediate ~f:(fun label -> Map.find_exn m label)) - immediate + Set.union (union_map immediate ~f:(fun label -> Map.find_exn m label)) immediate in Set.remove updated label in - let step_map m = - Map.mapi m ~f:(fun ~key:label ~data:_ -> step_node label m) - in + let step_map m = Map.mapi m ~f:(fun ~key:label ~data:_ -> step_node label m) in let map_equal = Map.Poly.equal Set.Poly.equal in let rec step_until_fixed m = let m' = step_map m in if map_equal m m' then m else step_until_fixed m' in step_until_fixed immediate_map +;; -let mir_reaching_definitions (mir : Program.Typed.t) (stmt : Stmt.Located.t) : - (label, reaching_defn Set.Poly.t entry_exit) Map.Poly.t = - let flowgraph, flowgraph_to_mir = - Monotone_framework.forward_flowgraph_of_stmt stmt - in +let mir_reaching_definitions (mir : Program.Typed.t) (stmt : Stmt.Located.t) + : (label, reaching_defn Set.Poly.t entry_exit) Map.Poly.t + = + let flowgraph, flowgraph_to_mir = Monotone_framework.forward_flowgraph_of_stmt stmt in let (module Flowgraph) = flowgraph in - let rd_map = - reaching_definitions_mfp mir (module Flowgraph) flowgraph_to_mir - in + let rd_map = reaching_definitions_mfp mir (module Flowgraph) flowgraph_to_mir in let to_rd_set set = - Set.Poly.map set ~f:(fun (s, label_opt) -> - (VVar s, Option.value label_opt ~default:1) ) + Set.Poly.map set ~f:(fun (s, label_opt) -> VVar s, Option.value label_opt ~default:1) in - Map.Poly.map rd_map ~f:(fun {entry; exit} -> - {entry= to_rd_set entry; exit= to_rd_set exit} ) + Map.Poly.map rd_map ~f:(fun { entry; exit } -> + { entry = to_rd_set entry; exit = to_rd_set exit }) +;; let all_labels - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) : int Set.Poly.t = + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + : int Set.Poly.t + = let step set = - Set.Poly.union set + Set.Poly.union + set (union_map set ~f:(fun l -> Map.Poly.find_exn Flowgraph.successors l)) in let rec step_fix set = @@ -141,19 +150,24 @@ let all_labels if Set.Poly.equal set next then set else step_fix next in step_fix Flowgraph.initials +;; let prog_rhs_variables (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (labels : int Set.Poly.t) : string Set.Poly.t = + (labels : int Set.Poly.t) + : string Set.Poly.t + = let label_vars label = Set.Poly.map ~f:(fun (VVar s, _) -> s) (stmt_rhs_var_set (Map.Poly.find_exn flowgraph_to_mir label).pattern) in union_map labels ~f:label_vars +;; -let stmt_uninitialized_variables (exceptions : string Set.Poly.t) - (stmt : Stmt.Located.t) : (Location_span.t * string) Set.Poly.t = +let stmt_uninitialized_variables (exceptions : string Set.Poly.t) (stmt : Stmt.Located.t) + : (Location_span.t * string) Set.Poly.t + = let flowgraph, flowgraph_to_mir = Monotone_framework.forward_flowgraph_of_stmt ~flatten_loops:true stmt in @@ -164,72 +178,76 @@ let stmt_uninitialized_variables (exceptions : string Set.Poly.t) initialized_vars_mfp all_variables (module Flowgraph) flowgraph_to_mir in let uninitialized = - Map.Poly.fold initialized_vars_map ~init:Set.Poly.empty + Map.Poly.fold + initialized_vars_map + ~init:Set.Poly.empty ~f:(fun ~key:label ~data:inits acc -> let stmt = Map.Poly.find_exn flowgraph_to_mir label in let rhs = Set.Poly.map - ~f:(fun (VVar s, {loc; _}) -> (loc, s)) + ~f:(fun (VVar s, { loc; _ }) -> loc, s) (stmt_rhs_var_set stmt.pattern) in let uninitialized (_, var) = not (Set.Poly.mem inits.entry var) in let uninitialized_set = Set.Poly.filter ~f:uninitialized rhs in - Set.Poly.union acc uninitialized_set ) + Set.Poly.union acc uninitialized_set) in - Set.Poly.filter uninitialized ~f:(fun (_, v) -> - not (Set.Poly.mem exceptions v) ) + Set.Poly.filter uninitialized ~f:(fun (_, v) -> not (Set.Poly.mem exceptions v)) +;; -let mir_uninitialized_variables (mir : Program.Typed.t) : - (Location_span.t * string) Set.Poly.t = +let mir_uninitialized_variables (mir : Program.Typed.t) + : (Location_span.t * string) Set.Poly.t + = let flag_variables = List.map ~f:Flag_vars.to_string Flag_vars.enumerate in let data_vars = data_set ~exclude_transformed:true mir in let trans_data_vars = data_set ~exclude_transformed:false mir in let globals = - Set.Poly.union - (Set.Poly.of_list flag_variables) - (Set.Poly.singleton "target") + Set.Poly.union (Set.Poly.of_list flag_variables) (Set.Poly.singleton "target") in let parameters = Set.Poly.of_list - (List.map ~f:fst + (List.map + ~f:fst (List.filter - ~f:(fun (_, {out_block; _}) -> out_block = Parameters) + ~f:(fun (_, { out_block; _ }) -> out_block = Parameters) mir.output_vars)) in let globals_data = Set.Poly.union globals data_vars in let globals_data_prep = - Set.Poly.union_list [globals_data; trans_data_vars; parameters] + Set.Poly.union_list [ globals_data; trans_data_vars; parameters ] in Set.Poly.union_list [ (* prepare_data scope: data *) - stmt_uninitialized_variables globals_data - {pattern= SList mir.prepare_data; meta= Location_span.empty} + stmt_uninitialized_variables + globals_data + { pattern = SList mir.prepare_data; meta = Location_span.empty } (* log_prob scope: data, prep declarations *) - ; stmt_uninitialized_variables globals_data_prep - {pattern= SList mir.log_prob; meta= Location_span.empty} + ; stmt_uninitialized_variables + globals_data_prep + { pattern = SList mir.log_prob; meta = Location_span.empty } (* gen quant scope: data, prep declarations *) - ; stmt_uninitialized_variables globals_data_prep - {pattern= SList mir.generate_quantities; meta= Location_span.empty} + ; stmt_uninitialized_variables + globals_data_prep + { pattern = SList mir.generate_quantities; meta = Location_span.empty } (* functions scope: arguments *) ; Set.Poly.union_list - (List.map mir.functions_block ~f:(fun {fdbody; fdargs; _} -> + (List.map mir.functions_block ~f:(fun { fdbody; fdargs; _ } -> let arg_vars = - Set.Poly.of_list - (List.map fdargs ~f:(fun (_, arg_name, _) -> arg_name)) + Set.Poly.of_list (List.map fdargs ~f:(fun (_, arg_name, _) -> arg_name)) in - stmt_uninitialized_variables - (Set.Poly.union arg_vars globals) - fdbody )) ] + stmt_uninitialized_variables (Set.Poly.union arg_vars globals) fdbody)) + ] +;; -let build_dep_info_map (mir : Program.Typed.t) - (stmt : (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t) : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t = +let build_dep_info_map + (mir : Program.Typed.t) + (stmt : (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t) + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t + = let statement_map = build_statement_map - (fun Stmt.Fixed.({pattern; _}) -> pattern) - (fun Stmt.Fixed.({meta; _}) -> meta) + (fun Stmt.Fixed.{ pattern; _ } -> pattern) + (fun Stmt.Fixed.{ meta; _ } -> meta) stmt in let _, preds, parents = build_cf_graphs statement_map in @@ -237,22 +255,26 @@ let build_dep_info_map (mir : Program.Typed.t) Map.Poly.mapi statement_map ~f:(fun ~key:label ~data:(stmt, meta) -> let rds = Map.find_exn rd_map label in ( stmt - , { predecessors= Map.find_exn preds label - ; parents= Map.find_exn parents label - ; reaching_defn_entry= rds.entry - ; reaching_defn_exit= rds.exit - ; meta } ) ) + , { predecessors = Map.find_exn preds label + ; parents = Map.find_exn parents label + ; reaching_defn_entry = rds.entry + ; reaching_defn_exit = rds.exit + ; meta + } )) +;; -let log_prob_build_dep_info_map (mir : Program.Typed.t) : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t = +let log_prob_build_dep_info_map (mir : Program.Typed.t) + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t + = let log_prob_stmt = - Stmt.Fixed.{meta= Location_span.empty; pattern= SList mir.log_prob} + Stmt.Fixed.{ meta = Location_span.empty; pattern = SList mir.log_prob } in build_dep_info_map mir log_prob_stmt +;; -let log_prob_dependency_graph (mir : Program.Typed.t) : - (label, label Set.Poly.t) Map.Poly.t = +let log_prob_dependency_graph (mir : Program.Typed.t) + : (label, label Set.Poly.t) Map.Poly.t + = let dep_info_map = log_prob_build_dep_info_map mir in all_node_dependencies dep_info_map +;; diff --git a/src/analysis_and_optimization/Dependence_analysis.mli b/src/analysis_and_optimization/Dependence_analysis.mli index 29b035a04c..f1a3ad3f69 100644 --- a/src/analysis_and_optimization/Dependence_analysis.mli +++ b/src/analysis_and_optimization/Dependence_analysis.mli @@ -22,94 +22,80 @@ open Dataflow_types Label dependence doesn't need the exit RD set, but variable dependence does. *) type node_dep_info = - { predecessors: label Set.Poly.t - ; parents: label Set.Poly.t - ; reaching_defn_entry: reaching_defn Set.Poly.t - ; reaching_defn_exit: reaching_defn Set.Poly.t - ; meta: Location_span.t } + { predecessors : label Set.Poly.t + ; parents : label Set.Poly.t + ; reaching_defn_entry : reaching_defn Set.Poly.t + ; reaching_defn_exit : reaching_defn Set.Poly.t + ; meta : Location_span.t + } -val node_immediate_dependencies : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t - -> ?blockers:vexpr Set.Poly.t - -> label - -> label Set.Poly.t (** Given dependency information for each node, find the 'immediate' dependencies of a node, where 'immediate' means the first-degree control flow parents and the reachable definitions of RHS variables. *) - -val node_dependencies : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t +val node_immediate_dependencies + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t + -> ?blockers:vexpr Set.Poly.t -> label -> label Set.Poly.t + (** Given dependency information for each node, find all of the dependencies of a single node. *) - -val node_vars_dependencies : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t - -> ?blockers:vexpr Set.Poly.t - -> vexpr Set.Poly.t +val node_dependencies + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t -> label -> label Set.Poly.t + (** Given dependency information for each node, find all of the dependencies of a set of variables at single node. 'blockers' are variables which will not be traversed. *) +val node_vars_dependencies + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t + -> ?blockers:vexpr Set.Poly.t + -> vexpr Set.Poly.t + -> label + -> label Set.Poly.t -val build_dep_info_map : - Program.Typed.t - -> (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t - -> ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t (** Build the dependency information for each node in the log_prob section of a program *) +val build_dep_info_map + : Program.Typed.t + -> (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t + -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t -val log_prob_build_dep_info_map : - Program.Typed.t - -> ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t (** Build the dependency information for each node in the log_prob section of a program *) +val log_prob_build_dep_info_map + : Program.Typed.t + -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t -val all_node_dependencies : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t - -> (label, label Set.Poly.t) Map.Poly.t (** Given dependency information for each node, find all of the dependencies of all nodes, effectively building the dependency graph. This is more efficient than calling node_dependencies on each node individually. *) +val all_node_dependencies + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t + -> (label, label Set.Poly.t) Map.Poly.t -val log_prob_dependency_graph : - Program.Typed.t -> (label, label Set.Poly.t) Map.Poly.t (** Build the dependency graph for the log_prob section of a program, where labels correspond to the labels built by statement_map. *) +val log_prob_dependency_graph : Program.Typed.t -> (label, label Set.Poly.t) Map.Poly.t -val reaching_defn_lookup : - reaching_defn Set.Poly.t -> vexpr -> label Set.Poly.t +val reaching_defn_lookup : reaching_defn Set.Poly.t -> vexpr -> label Set.Poly.t -val mir_uninitialized_variables : - Program.Typed.t -> (Location_span.t * string) Set.Poly.t (** Produce a list of uninitialized variables and their label locations, from the flowgraph starting at the given statement *) +val mir_uninitialized_variables : Program.Typed.t -> (Location_span.t * string) Set.Poly.t diff --git a/src/analysis_and_optimization/Factor_graph.ml b/src/analysis_and_optimization/Factor_graph.ml index 57dd910986..30a36ba4b0 100644 --- a/src/analysis_and_optimization/Factor_graph.ml +++ b/src/analysis_and_optimization/Factor_graph.ml @@ -14,97 +14,96 @@ type factor = [@@deriving sexp, hash, compare] type factor_graph = - { factor_map: (factor * label, vexpr Set.Poly.t) Map.Poly.t - ; var_map: (vexpr, (factor * label) Set.Poly.t) Map.Poly.t } + { factor_map : (factor * label, vexpr Set.Poly.t) Map.Poly.t + ; var_map : (vexpr, (factor * label) Set.Poly.t) Map.Poly.t + } [@@deriving sexp, compare] let extract_factors_statement stmt = match stmt with | Stmt.Fixed.Pattern.TargetPE e -> - List.map (summation_terms e) ~f:(fun x -> TargetTerm x) - | NRFunApp (_, f, _) when Internal_fun.of_string_opt f = Some FnReject -> - [Reject] - | NRFunApp (_, s, args) when String.suffix s 3 = "_lp" -> - [LPFunction (s, args)] + List.map (summation_terms e) ~f:(fun x -> TargetTerm x) + | NRFunApp (_, f, _) when Internal_fun.of_string_opt f = Some FnReject -> [ Reject ] + | NRFunApp (_, s, args) when String.suffix s 3 = "_lp" -> [ LPFunction (s, args) ] | Assignment (_, _) - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |For _ | Block _ | SList _ - |Decl {decl_id= _; _} -> - [] + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | For _ | Block _ | SList _ + | Decl { decl_id = _; _ } -> [] +;; let rec extract_factors statement_map label = let stmt, _ = Map.Poly.find_exn statement_map label in - let this_stmt = - List.map (extract_factors_statement stmt) ~f:(fun x -> (label, x)) - in + let this_stmt = List.map (extract_factors_statement stmt) ~f:(fun x -> label, x) in Stmt.Fixed.Pattern.fold (fun s _ -> s) (fun state label -> List.append state (extract_factors statement_map label)) - this_stmt stmt + this_stmt + stmt +;; let factor_rhs (factor : factor) : vexpr Set.Poly.t = match factor with | TargetTerm e -> Set.Poly.map (expr_var_set e) ~f:fst | Reject -> Set.Poly.empty | LPFunction (_, es) -> Set.Poly.of_list (List.map es ~f:vexpr_of_expr_exn) +;; let factor_var_dependencies statement_map blockers (label, factor) = let rhs = factor_rhs factor in let dep_labels = node_vars_dependencies statement_map ~blockers rhs label in let label_vars l = - Set.Poly.map - (stmt_rhs_var_set (fst (Map.Poly.find_exn statement_map l))) - ~f:fst + Set.Poly.map (stmt_rhs_var_set (fst (Map.Poly.find_exn statement_map l))) ~f:fst in let dep_vars = union_map dep_labels ~f:label_vars in Set.Poly.union dep_vars rhs +;; (* Helper function to generate the factor graph adjacency map representation from a factor-adjacency list *) let build_adjacency_maps (factors : (label * factor * vexpr Set.Poly.t) List.t) - : factor_graph = + : factor_graph + = let factor_map = - List.fold ~f:merge_set_maps ~init:Map.Poly.empty - (List.map - ~f:(fun (l, fac, vars) -> Map.Poly.singleton (fac, l) vars) - factors) + List.fold + ~f:merge_set_maps + ~init:Map.Poly.empty + (List.map ~f:(fun (l, fac, vars) -> Map.Poly.singleton (fac, l) vars) factors) in let var_map = - List.fold ~f:merge_set_maps ~init:Map.Poly.empty + List.fold + ~f:merge_set_maps + ~init:Map.Poly.empty (List.concat_map factors ~f:(fun (l, fac, vars) -> List.map ~f:(fun v -> Map.Poly.singleton v (Set.Poly.singleton (fac, l))) - (Set.Poly.to_list vars) )) + (Set.Poly.to_list vars))) in - {factor_map; var_map} + { factor_map; var_map } +;; -let fg_remove_fac (fac : factor * cf_state) (fg : factor_graph) : factor_graph - = +let fg_remove_fac (fac : factor * cf_state) (fg : factor_graph) : factor_graph = let factor_map = Map.Poly.remove fg.factor_map fac in - {fg with factor_map} + { fg with factor_map } +;; let fg_remove_var (var : vexpr) (fg : factor_graph) : factor_graph = - let factor_map = - Map.Poly.map fg.factor_map ~f:(fun vars -> Set.Poly.remove vars var) - in + let factor_map = Map.Poly.map fg.factor_map ~f:(fun vars -> Set.Poly.remove vars var) in let var_map = Map.Poly.remove fg.var_map var in - {factor_map; var_map} + { factor_map; var_map } +;; let remove_touching vars fg = let facs = union_map vars ~f:(fun v -> - Option.value ~default:Set.Poly.empty (Map.Poly.find fg.var_map v) ) - in - let without_vars = - Set.fold ~f:(fun g v -> fg_remove_var v g) ~init:fg vars - in - let without_facs = - Set.fold ~f:(fun g f -> fg_remove_fac f g) ~init:without_vars facs + Option.value ~default:Set.Poly.empty (Map.Poly.find fg.var_map v)) in + let without_vars = Set.fold ~f:(fun g v -> fg_remove_var v g) ~init:fg vars in + let without_facs = Set.fold ~f:(fun g f -> fg_remove_fac f g) ~init:without_vars facs in without_facs +;; (* Build a factor graph from prog.log_prob using dependency analysis *) let prog_factor_graph ?(exclude_data_facs : bool = false) prog : factor_graph = @@ -118,20 +117,19 @@ let prog_factor_graph ?(exclude_data_facs : bool = false) prog : factor_graph = in let factor_list = List.map factors ~f:(fun (l, fac) -> - ( l - , fac - , Set.Poly.inter vars - (factor_var_dependencies statement_map vars (l, fac)) ) ) + l, fac, Set.Poly.inter vars (factor_var_dependencies statement_map vars (l, fac))) in let fg = build_adjacency_maps factor_list in - if exclude_data_facs then - remove_touching (Set.Poly.map ~f:(fun v -> VVar v) data_vars) fg + if exclude_data_facs + then remove_touching (Set.Poly.map ~f:(fun v -> VVar v) data_vars) fg else fg +;; (* BFS on 'fg' with initial frontier 'starts' and terminating at any element of 'goals' *) -let fg_reaches (starts : vexpr Set.Poly.t) (goals : vexpr Set.Poly.t) - (fg : factor_graph) : bool = +let fg_reaches (starts : vexpr Set.Poly.t) (goals : vexpr Set.Poly.t) (fg : factor_graph) + : bool + = let vneighbors v = let factors = Map.Poly.find_exn fg.var_map v in union_map factors ~f:(Map.Poly.find_exn fg.factor_map) @@ -139,28 +137,41 @@ let fg_reaches (starts : vexpr Set.Poly.t) (goals : vexpr Set.Poly.t) let rec step (frontier : vexpr List.t) (visited : vexpr Set.Poly.t) = match frontier with | next :: frontier' -> - if Set.mem visited next then step frontier' visited - else - let visited' = Set.Poly.add visited next in - let expansion = vneighbors next in - if not (Set.Poly.is_empty (Set.Poly.inter expansion goals)) then true - else - step (List.append frontier' (Set.Poly.to_list expansion)) visited' + if Set.mem visited next + then step frontier' visited + else ( + let visited' = Set.Poly.add visited next in + let expansion = vneighbors next in + if not (Set.Poly.is_empty (Set.Poly.inter expansion goals)) + then true + else step (List.append frontier' (Set.Poly.to_list expansion)) visited') | [] -> false in step (Set.Poly.to_list starts) Set.Poly.empty - -let fg_factor_reaches (start : factor * label) (goals : vexpr Set.Poly.t) - (fg : factor_graph) : bool = +;; + +let fg_factor_reaches + (start : factor * label) + (goals : vexpr Set.Poly.t) + (fg : factor_graph) + : bool + = let var_starts = Map.Poly.find_exn fg.factor_map start in fg_reaches var_starts goals fg - -let fg_factor_is_prior (var : vexpr) (fac : factor * label) - (data : vexpr Set.Poly.t) (fg : factor_graph) : bool = +;; + +let fg_factor_is_prior + (var : vexpr) + (fac : factor * label) + (data : vexpr Set.Poly.t) + (fg : factor_graph) + : bool + = (* build G'=G\V *) let fg' = fg_remove_var var fg in (* Check if the data is now unreachable *) not (fg_factor_reaches fac data fg') +;; (* Priors of V are neighbors of V which have no connection to any data except though V So for graph G and each parameter V: @@ -169,39 +180,47 @@ let fg_factor_is_prior (var : vexpr) (fac : factor * label) Use BFS starting from F in G' and search for any data, if there is none, F is a prior *) -let fg_var_priors (var : vexpr) (data : vexpr Set.Poly.t) (fg : factor_graph) : - (factor * label) Set.Poly.t option = +let fg_var_priors (var : vexpr) (data : vexpr Set.Poly.t) (fg : factor_graph) + : (factor * label) Set.Poly.t option + = match Map.Poly.find fg.var_map var with | Some factors -> - Some - (Set.Poly.filter factors ~f:(fun fac -> - fg_factor_is_prior var fac data fg )) + Some (Set.Poly.filter factors ~f:(fun fac -> fg_factor_is_prior var fac data fg)) | None -> None +;; -let list_priors ?factor_graph:(fg_opt = None) (mir : Program.Typed.t) : - (vexpr, (factor * label) Set.Poly.t option) Map.Poly.t = +let list_priors ?factor_graph:(fg_opt = None) (mir : Program.Typed.t) + : (vexpr, (factor * label) Set.Poly.t option) Map.Poly.t + = let fg = Option.value ~default:(prog_factor_graph mir) fg_opt in let params = Set.Poly.map ~f:(fun v -> VVar v) (parameter_names_set mir) in let data = Set.Poly.map ~f:(fun v -> VVar v) (data_set mir) in let likely_sizes = - Set.Poly.diff data + Set.Poly.diff + data (Set.Poly.map ~f:(fun v -> VVar v) (data_set ~exclude_ints:true mir)) in let fg' = - Set.Poly.fold ~init:fg + Set.Poly.fold + ~init:fg ~f:(fun fg likely_size -> fg_remove_var likely_size fg) likely_sizes in (* for each param, apply fg_var_priors and collect results in a map*) generate_map params ~f:(fun p -> fg_var_priors p data fg') +;; let string_of_factor (factor : factor) : string = match factor with | TargetTerm e -> Fmt.strf "\"%a\"" Expr.Typed.pp e | Reject -> "reject" | LPFunction (s, _) -> s +;; -let string_of_vexpr (vexpr : vexpr) : string = match vexpr with VVar s -> s +let string_of_vexpr (vexpr : vexpr) : string = + match vexpr with + | VVar s -> s +;; (* Utility to print a factor graph to the Graphviz dot language for visualization *) @@ -210,21 +229,22 @@ let factor_graph_to_dot (fg : factor_graph) : string = let names = List.map ~f:(fun ((f, _), ps) -> - (string_of_factor f, List.map ~f:string_of_vexpr (Set.Poly.to_list ps)) - ) + string_of_factor f, List.map ~f:string_of_vexpr (Set.Poly.to_list ps)) factors in let factor_names, param_name_lists = List.unzip names in let factor_strings = - List.map factor_names ~f:(fun n -> String.concat [n; " [shape=box]"]) + List.map factor_names ~f:(fun n -> String.concat [ n; " [shape=box]" ]) in let param_strings = List.dedup_and_sort ~compare:String.compare (List.concat param_name_lists) in let edge_strings = List.concat_map - ~f:(fun (f, ps) -> List.map ~f:(fun p -> String.concat [f; " -- "; p]) ps) + ~f:(fun (f, ps) -> List.map ~f:(fun p -> String.concat [ f; " -- "; p ]) ps) names in - [["graph {"]; factor_strings; param_strings; edge_strings; ["}"]] - |> List.concat |> String.concat ~sep:"\n" + [ [ "graph {" ]; factor_strings; param_strings; edge_strings; [ "}" ] ] + |> List.concat + |> String.concat ~sep:"\n" +;; diff --git a/src/analysis_and_optimization/Mir_utils.ml b/src/analysis_and_optimization/Mir_utils.ml index 636199805e..a0f78de6f1 100644 --- a/src/analysis_and_optimization/Mir_utils.ml +++ b/src/analysis_and_optimization/Mir_utils.ml @@ -7,10 +7,11 @@ open Dataflow_types let rec fold_expr ~take_expr ~(init : 'c) (expr : Expr.Typed.t) : 'c = Expr.Fixed.Pattern.fold_left ~f:(fun a e -> fold_expr ~take_expr ~init:(take_expr a e) e) - ~init expr.pattern + ~init + expr.pattern +;; -let fold_stmts ~take_expr ~take_stmt ~(init : 'c) - (stmts : Stmt.Located.t List.t) : 'c = +let fold_stmts ~take_expr ~take_stmt ~(init : 'c) (stmts : Stmt.Located.t List.t) : 'c = (* let rec fold_expr (state : 'c) (expr : Expr.Typed.Meta.t Expr.Fixed.t) = * Expr.Fixed.Pattern.fold_left * ~f:(fun a e -> fold_expr (take_expr a e) e) @@ -21,297 +22,331 @@ let fold_stmts ~take_expr ~take_stmt ~(init : 'c) Stmt.Fixed.Pattern.fold_left ~f:(fun a e -> fold_expr ~take_expr ~init:(take_expr a e) e) ~g:(fun a s -> fold_stmt (take_stmt a s) s) - ~init:state stmt.pattern + ~init:state + stmt.pattern in List.fold ~f:(fun a s -> fold_stmt (take_stmt a s) s) ~init stmts +;; let rec num_expr_value (v : Expr.Typed.t) : (float * string) option = match v with - | {pattern= Fixed.Pattern.Lit (Real, str); _} - |{pattern= Fixed.Pattern.Lit (Int, str); _} -> - Some (float_of_string str, str) - | {pattern= Fixed.Pattern.FunApp (StanLib, "PMinus__", [v]); _} -> ( - match num_expr_value v with + | { pattern = Fixed.Pattern.Lit (Real, str); _ } + | { pattern = Fixed.Pattern.Lit (Int, str); _ } -> Some (float_of_string str, str) + | { pattern = Fixed.Pattern.FunApp (StanLib, "PMinus__", [ v ]); _ } -> + (match num_expr_value v with | Some (v, s) -> Some (-.v, "-" ^ s) - | None -> None ) + | None -> None) | _ -> None +;; type bound_values = - { lower: [`None | `Nonlit | `Lit of float] - ; upper: [`None | `Nonlit | `Lit of float] } + { lower : [ `None | `Nonlit | `Lit of float ] + ; upper : [ `None | `Nonlit | `Lit of float ] + } let trans_bounds_values (trans : Expr.Typed.t transformation) : bound_values = let bound_value e = - match num_expr_value e with None -> `Nonlit | Some (f, _) -> `Lit f + match num_expr_value e with + | None -> `Nonlit + | Some (f, _) -> `Lit f in match trans with - | Lower lower -> {lower= bound_value lower; upper= `None} - | Upper upper -> {lower= `None; upper= bound_value upper} - | LowerUpper (lower, upper) -> - {lower= bound_value lower; upper= bound_value upper} - | Simplex -> {lower= `Lit 0.; upper= `Lit 1.} - | PositiveOrdered -> {lower= `Lit 0.; upper= `None} - | UnitVector -> {lower= `Lit (-1.); upper= `Lit 1.} - | CholeskyCorr | CholeskyCov | Correlation | Covariance | Ordered - |Offset _ | Multiplier _ | OffsetMultiplier _ | Identity -> - {lower= `None; upper= `None} + | Lower lower -> { lower = bound_value lower; upper = `None } + | Upper upper -> { lower = `None; upper = bound_value upper } + | LowerUpper (lower, upper) -> { lower = bound_value lower; upper = bound_value upper } + | Simplex -> { lower = `Lit 0.; upper = `Lit 1. } + | PositiveOrdered -> { lower = `Lit 0.; upper = `None } + | UnitVector -> { lower = `Lit (-1.); upper = `Lit 1. } + | CholeskyCorr + | CholeskyCov + | Correlation + | Covariance + | Ordered + | Offset _ + | Multiplier _ + | OffsetMultiplier _ + | Identity -> { lower = `None; upper = `None } +;; let chop_dist_name (fname : string) : string Option.t = (* Slightly inefficient, would be better to short-circuit *) - List.fold ~init:None ~f:Option.first_some + List.fold + ~init:None + ~f:Option.first_some (List.map ~f:(fun suffix -> String.chop_suffix ~suffix fname) - ["_propto_log"; "_propto_lpdf"; "_propto_lpmf"]) + [ "_propto_log"; "_propto_lpdf"; "_propto_lpmf" ]) +;; let is_dist (fname : string) : bool = Option.is_some (chop_dist_name fname) -let rec top_var_declarations Stmt.Fixed.({pattern; _}) : string Set.Poly.t = +let rec top_var_declarations Stmt.Fixed.{ pattern; _ } : string Set.Poly.t = match pattern with - | Decl {decl_id; _} -> Set.Poly.singleton decl_id + | Decl { decl_id; _ } -> Set.Poly.singleton decl_id | SList l -> Set.Poly.union_list (List.map ~f:top_var_declarations l) | _ -> Set.Poly.empty - -let data_set ?(exclude_transformed = false) ?(exclude_ints = false) - (mir : Program.Typed.t) : string Set.Poly.t = +;; + +let data_set + ?(exclude_transformed = false) + ?(exclude_ints = false) + (mir : Program.Typed.t) + : string Set.Poly.t + = (* Data are input_vars *) let data = Set.Poly.of_list mir.input_vars in (* Possibly remove ints from the data set *) let filtered_data = - let remove_ints = - Set.Poly.filter ~f:(fun (_, st) -> st <> SizedType.SInt) - in + let remove_ints = Set.Poly.filter ~f:(fun (_, st) -> st <> SizedType.SInt) in Set.Poly.map ~f:fst ((if exclude_ints then remove_ints else ident) data) in (* Transformed data are declarations in prepare_data but excluding data *) - if exclude_transformed then filtered_data - else + if exclude_transformed + then filtered_data + else ( let trans_data = Set.Poly.diff - (Set.Poly.union_list - (List.map ~f:top_var_declarations mir.prepare_data)) + (Set.Poly.union_list (List.map ~f:top_var_declarations mir.prepare_data)) (Set.Poly.map ~f:fst data) in - Set.Poly.union trans_data filtered_data + Set.Poly.union trans_data filtered_data) +;; let parameter_set ?(include_transformed = false) (mir : Program.Typed.t) = Set.Poly.of_list (List.map - ~f:(fun (pname, {out_trans; _}) -> (pname, out_trans)) + ~f:(fun (pname, { out_trans; _ }) -> pname, out_trans) (List.filter - ~f:(fun (_, {out_block; _}) -> + ~f:(fun (_, { out_block; _ }) -> out_block = Parameters - || (include_transformed && out_block = TransformedParameters) ) + || (include_transformed && out_block = TransformedParameters)) mir.output_vars)) +;; -let parameter_names_set ?(include_transformed = false) (mir : Program.Typed.t) - = +let parameter_names_set ?(include_transformed = false) (mir : Program.Typed.t) = Set.Poly.map ~f:fst (parameter_set ~include_transformed mir) +;; -let rec var_declarations Stmt.Fixed.({pattern; _}) : string Set.Poly.t = +let rec var_declarations Stmt.Fixed.{ pattern; _ } : string Set.Poly.t = match pattern with - | Decl {decl_id; _} -> Set.Poly.singleton decl_id - | IfElse (_, s, None) | While (_, s) | For {body= s; _} -> var_declarations s - | IfElse (_, s1, Some s2) -> - Set.Poly.union (var_declarations s1) (var_declarations s2) - | Block slist | SList slist -> - Set.Poly.union_list (List.map ~f:var_declarations slist) + | Decl { decl_id; _ } -> Set.Poly.singleton decl_id + | IfElse (_, s, None) | While (_, s) | For { body = s; _ } -> var_declarations s + | IfElse (_, s1, Some s2) -> Set.Poly.union (var_declarations s1) (var_declarations s2) + | Block slist | SList slist -> Set.Poly.union_list (List.map ~f:var_declarations slist) | _ -> Set.Poly.empty +;; let rec map_rec_expr f e = let recurse = map_rec_expr f in - Expr.Fixed.{e with pattern= f (Pattern.map recurse e.pattern)} + Expr.Fixed.{ e with pattern = f (Pattern.map recurse e.pattern) } +;; let map_rec_expr_state f state e = let cur_state = ref state in let g e' = let e', state = f !cur_state e' in - cur_state := state ; + cur_state := state; e' in let e = map_rec_expr g e in let state = !cur_state in - (e, state) + e, state +;; let rec map_rec_stmt_loc f stmt = let recurse = map_rec_stmt_loc f in - Stmt.Fixed. - {stmt with pattern= f (Pattern.map (fun x -> x) recurse stmt.pattern)} + Stmt.Fixed.{ stmt with pattern = f (Pattern.map (fun x -> x) recurse stmt.pattern) } +;; let rec top_down_map_rec_stmt_loc f stmt = let recurse = top_down_map_rec_stmt_loc f in - Stmt.Fixed.{stmt with pattern= Pattern.map Fn.id recurse (f stmt.pattern)} + Stmt.Fixed.{ stmt with pattern = Pattern.map Fn.id recurse (f stmt.pattern) } +;; let map_rec_state_stmt_loc f state stmt = let cur_state = ref state in let g stmt = let stmt, state = f !cur_state stmt in - cur_state := state ; + cur_state := state; stmt in let stmt = map_rec_stmt_loc g stmt in let state = !cur_state in - (stmt, state) + stmt, state +;; let map_rec_stmt_loc_num flowgraph_to_mir f s = - let rec map_rec_stmt_loc_num' (cur_node : int) - (stmt : Stmt.Located.Non_recursive.t) = + let rec map_rec_stmt_loc_num' (cur_node : int) (stmt : Stmt.Located.Non_recursive.t) = let find_node i = Map.find_exn flowgraph_to_mir i in let recurse i = map_rec_stmt_loc_num' i (find_node i) in Stmt.Fixed. - { pattern= f cur_node (Pattern.map Fn.id recurse stmt.pattern) - ; meta= stmt.meta } + { pattern = f cur_node (Pattern.map Fn.id recurse stmt.pattern); meta = stmt.meta } in map_rec_stmt_loc_num' 1 s +;; let map_rec_state_stmt_loc_num (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) (f : - int + int -> 's -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t - -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t * 's) (state : 's) - (s : Stmt.Located.Non_recursive.t) : Stmt.Located.t * 's = + -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t * 's) + (state : 's) + (s : Stmt.Located.Non_recursive.t) + : Stmt.Located.t * 's + = let cur_state = ref state in let g i stmt = let stmt, state = f i !cur_state stmt in - cur_state := state ; + cur_state := state; stmt in let stmt = map_rec_stmt_loc_num flowgraph_to_mir g s in let state = !cur_state in - (stmt, state) + stmt, state +;; let stmt_loc_of_stmt_loc_num flowgraph_to_mir s = (* (flowgraph_to_mir : (int, stmt_loc_num) Map.Poly.t) (s : stmt_loc_num) = *) map_rec_stmt_loc_num flowgraph_to_mir (fun _ s' -> s') s +;; let statement_stmt_loc_of_statement_stmt_loc_num flowgraph_to_mir pattern = - (stmt_loc_of_stmt_loc_num flowgraph_to_mir - Stmt.Located.{Non_recursive.pattern; meta= Meta.empty}) + (stmt_loc_of_stmt_loc_num + flowgraph_to_mir + Stmt.Located.{ Non_recursive.pattern; meta = Meta.empty }) .pattern +;; (** Forgetful function from numbered to unnumbered programs *) let unnumbered_prog_of_numbered_prog flowgraph_to_mir p = Program.map (stmt_loc_of_stmt_loc_num flowgraph_to_mir) p +;; (** See interface file *) let fwd_traverse_statement stmt ~init ~f = Stmt.Fixed.Pattern.( match stmt with | IfElse (pred, then_s, else_s_opt) -> - let s', c = f init then_s in - Option.value_map else_s_opt - ~default:(s', IfElse (pred, c, None)) - ~f:(fun else_s -> - let s'', c' = f s' else_s in - (s'', IfElse (pred, c, Some c')) ) + let s', c = f init then_s in + Option.value_map + else_s_opt + ~default:(s', IfElse (pred, c, None)) + ~f:(fun else_s -> + let s'', c' = f s' else_s in + s'', IfElse (pred, c, Some c')) | While (pred, body) -> - let s', c = f init body in - (s', While (pred, c)) + let s', c = f init body in + s', While (pred, c) | For vars -> - let s', c = f init vars.body in - (s', For {vars with body= c}) + let s', c = f init vars.body in + s', For { vars with body = c } | Block stmts -> - let s', ls = - List.fold_left stmts - ~f:(fun (s, l) stmt -> - let s', c = f s stmt in - (s', List.cons c l) ) - ~init:(init, []) - in - (s', Block (List.rev ls)) + let s', ls = + List.fold_left + stmts + ~f:(fun (s, l) stmt -> + let s', c = f s stmt in + s', List.cons c l) + ~init:(init, []) + in + s', Block (List.rev ls) | SList stmts -> - let s', ls = - List.fold_left stmts - ~f:(fun (s, l) stmt -> - let s', c = f s stmt in - (s', List.cons c l) ) - ~init:(init, []) - in - (s', SList (List.rev ls)) - | Assignment _ as s -> (init, s) - | TargetPE _ as s -> (init, s) - | NRFunApp _ as s -> (init, s) - | Break as s -> (init, s) - | Continue as s -> (init, s) - | Return _ as s -> (init, s) - | Skip as s -> (init, s) - | Decl _ as s -> (init, s)) + let s', ls = + List.fold_left + stmts + ~f:(fun (s, l) stmt -> + let s', c = f s stmt in + s', List.cons c l) + ~init:(init, []) + in + s', SList (List.rev ls) + | Assignment _ as s -> init, s + | TargetPE _ as s -> init, s + | NRFunApp _ as s -> init, s + | Break as s -> init, s + | Continue as s -> init, s + | Return _ as s -> init, s + | Skip as s -> init, s + | Decl _ as s -> init, s) +;; (** See interface file *) -let vexpr_of_expr_exn Expr.Fixed.({pattern; _}) = +let vexpr_of_expr_exn Expr.Fixed.{ pattern; _ } = match pattern with | Var s -> VVar s | _ -> raise (Failure "Non-var expression found, but var expected") +;; (** See interface file *) -let rec expr_var_set Expr.Fixed.({pattern; meta}) = - let union_recur exprs = - Set.Poly.union_list (List.map exprs ~f:expr_var_set) - in +let rec expr_var_set Expr.Fixed.{ pattern; meta } = + let union_recur exprs = Set.Poly.union_list (List.map exprs ~f:expr_var_set) in match pattern with | Var s -> Set.Poly.singleton (VVar s, meta) | Lit _ -> Set.Poly.empty | FunApp (_, _, exprs) -> union_recur exprs - | TernaryIf (expr1, expr2, expr3) -> union_recur [expr1; expr2; expr3] + | TernaryIf (expr1, expr2, expr3) -> union_recur [ expr1; expr2; expr3 ] | Indexed (expr, ix) -> - Set.Poly.union_list (expr_var_set expr :: List.map ix ~f:index_var_set) - | EAnd (expr1, expr2) | EOr (expr1, expr2) -> union_recur [expr1; expr2] + Set.Poly.union_list (expr_var_set expr :: List.map ix ~f:index_var_set) + | EAnd (expr1, expr2) | EOr (expr1, expr2) -> union_recur [ expr1; expr2 ] and index_var_set ix = match ix with | All -> Set.Poly.empty | Single expr -> expr_var_set expr | Upfrom expr -> expr_var_set expr - | Between (expr1, expr2) -> - Set.Poly.union (expr_var_set expr1) (expr_var_set expr2) + | Between (expr1, expr2) -> Set.Poly.union (expr_var_set expr1) (expr_var_set expr2) | MultiIndex expr -> expr_var_set expr +;; let stmt_rhs stmt = match stmt with - | Stmt.Fixed.Pattern.For vars -> Set.Poly.of_list [vars.lower; vars.upper] + | Stmt.Fixed.Pattern.For vars -> Set.Poly.of_list [ vars.lower; vars.upper ] | NRFunApp (_, _, exprs) -> Set.Poly.of_list exprs | IfElse (rhs, _, _) - |While (rhs, _) - |Assignment (_, rhs) - |TargetPE rhs - |Return (Some rhs) -> - Set.Poly.singleton rhs - | Return None | Break | Continue | Skip | Decl _ | Block _ | SList _ -> - Set.Poly.empty + | While (rhs, _) + | Assignment (_, rhs) + | TargetPE rhs + | Return (Some rhs) -> Set.Poly.singleton rhs + | Return None | Break | Continue | Skip | Decl _ | Block _ | SList _ -> Set.Poly.empty +;; let union_map (set : ('a, 'c) Set_intf.Set.t) ~(f : 'a -> 'b Set.Poly.t) = Set.fold set ~init:Set.Poly.empty ~f:(fun s a -> Set.Poly.union s (f a)) +;; let stmt_rhs_var_set stmt = union_map (stmt_rhs stmt) ~f:expr_var_set (** See interface file *) -let expr_assigned_var Expr.Fixed.({pattern; _}) = +let expr_assigned_var Expr.Fixed.{ pattern; _ } = match pattern with | Var s -> VVar s - | Indexed ({pattern= Var s; _}, _) -> VVar s + | Indexed ({ pattern = Var s; _ }, _) -> VVar s | _ -> raise (Failure "Unimplemented: analysis of assigning to non-var") +;; (** See interface file *) -let rec summation_terms (Expr.Fixed.({pattern; _}) as rhs) = +let rec summation_terms (Expr.Fixed.{ pattern; _ } as rhs) = match pattern with - | FunApp (_, "Plus__", [e1; e2]) -> - List.append (summation_terms e1) (summation_terms e2) - | _ -> [rhs] + | FunApp (_, "Plus__", [ e1; e2 ]) -> + List.append (summation_terms e1) (summation_terms e2) + | _ -> [ rhs ] +;; (** See interface file *) -let stmt_of_block b = - Stmt.Fixed.{pattern= SList b; meta= Stmt.Located.Meta.empty} +let stmt_of_block b = Stmt.Fixed.{ pattern = SList b; meta = Stmt.Located.Meta.empty } let rec fn_subst_expr m e = match m e with | Some e' -> - (* let print_expr (e:Expr.Typed.t) = *) - (* [%sexp (e.pattern : Expr.Typed.Meta.t Expr.Fixed.t Expr.Fixed.Pattern.t)] |> Sexp.to_string *) - (* in *) - (* let _ = print_endline ("Replaced expr: " ^ print_expr e ^ " -> " ^ print_expr e') in *) - e' - | _ -> Expr.Fixed.{e with pattern= Pattern.map (fn_subst_expr m) e.pattern} + (* let print_expr (e:Expr.Typed.t) = *) + (* [%sexp (e.pattern : Expr.Typed.Meta.t Expr.Fixed.t Expr.Fixed.Pattern.t)] |> Sexp.to_string *) + (* in *) + (* let _ = print_endline ("Replaced expr: " ^ print_expr e ^ " -> " ^ print_expr e') in *) + e' + | _ -> Expr.Fixed.{ e with pattern = Pattern.map (fn_subst_expr m) e.pattern } +;; let fn_subst_idx m = Index.map (fn_subst_expr m) @@ -320,24 +355,27 @@ let fn_subst_stmt_base_helper g h b = match b with | Assignment ((x, ut, l), e2) -> Assignment ((x, ut, List.map ~f:h l), g e2) | x -> map g (fun y -> y) x) +;; -let fn_subst_stmt_base m = - fn_subst_stmt_base_helper (fn_subst_expr m) (fn_subst_idx m) - +let fn_subst_stmt_base m = fn_subst_stmt_base_helper (fn_subst_expr m) (fn_subst_idx m) let fn_subst_stmt m = map_rec_stmt_loc (fn_subst_stmt_base m) let name_map m (e : Expr.Typed.t) = match e.pattern with - | Var s -> ( - match Map.Poly.find m s with - | Some s' -> Some {e with pattern= Var s'} - | None -> None ) + | Var s -> + (match Map.Poly.find m s with + | Some s' -> Some { e with pattern = Var s' } + | None -> None) | _ -> None +;; let name_subst_stmt m = fn_subst_stmt (name_map m) let var_map m (e : Expr.Typed.t) = - match e.pattern with Var s -> Map.find m s | _ -> None + match e.pattern with + | Var s -> Map.find m s + | _ -> None +;; let subst_expr m e = fn_subst_expr (var_map m) e let subst_idx m = Index.map (subst_expr m) @@ -349,82 +387,87 @@ let expr_subst_idx m = Index.map (expr_subst_expr m) let expr_subst_stmt_base m = fn_subst_stmt_base_helper (expr_subst_expr m) (expr_subst_idx m) +;; let expr_subst_stmt m = map_rec_stmt_loc (expr_subst_stmt_base m) -let rec expr_depth Expr.Fixed.({pattern; _}) = +let rec expr_depth Expr.Fixed.{ pattern; _ } = match pattern with | Var _ | Lit (_, _) -> 0 | FunApp (_, _, l) -> - 1 - + Option.value ~default:0 - (List.max_elt ~compare:compare_int (List.map ~f:expr_depth l)) + 1 + + Option.value + ~default:0 + (List.max_elt ~compare:compare_int (List.map ~f:expr_depth l)) | TernaryIf (e1, e2, e3) -> - 1 - + Option.value ~default:0 - (List.max_elt ~compare:compare_int - (List.map ~f:expr_depth [e1; e2; e3])) + 1 + + Option.value + ~default:0 + (List.max_elt ~compare:compare_int (List.map ~f:expr_depth [ e1; e2; e3 ])) | Indexed (e, l) -> - 1 - + max (expr_depth e) - (Option.value ~default:0 - (List.max_elt ~compare:compare_int (List.map ~f:idx_depth l))) + 1 + + max + (expr_depth e) + (Option.value + ~default:0 + (List.max_elt ~compare:compare_int (List.map ~f:idx_depth l))) | EAnd (e1, e2) | EOr (e1, e2) -> - 1 - + Option.value ~default:0 - (List.max_elt ~compare:compare_int (List.map ~f:expr_depth [e1; e2])) + 1 + + Option.value + ~default:0 + (List.max_elt ~compare:compare_int (List.map ~f:expr_depth [ e1; e2 ])) and idx_depth i = match i with | All -> 0 | Single e | Upfrom e | MultiIndex e -> expr_depth e | Between (e1, e2) -> max (expr_depth e1) (expr_depth e2) +;; let ad_level_sup l = - if List.exists l ~f:(fun x -> Expr.Typed.adlevel_of x = AutoDiffable) then - UnsizedType.AutoDiffable + if List.exists l ~f:(fun x -> Expr.Typed.adlevel_of x = AutoDiffable) + then UnsizedType.AutoDiffable else DataOnly +;; -let rec update_expr_ad_levels autodiffable_variables - (Expr.Fixed.({pattern; _}) as e) = +let rec update_expr_ad_levels autodiffable_variables (Expr.Fixed.{ pattern; _ } as e) = match pattern with | Var x -> - if Set.Poly.mem autodiffable_variables x then - Expr.Typed.{e with meta= Meta.{e.meta with adlevel= AutoDiffable}} - else {e with meta= {e.meta with adlevel= DataOnly}} - | Lit (_, _) -> {e with meta= {e.meta with adlevel= DataOnly}} + if Set.Poly.mem autodiffable_variables x + then Expr.Typed.{ e with meta = Meta.{ e.meta with adlevel = AutoDiffable } } + else { e with meta = { e.meta with adlevel = DataOnly } } + | Lit (_, _) -> { e with meta = { e.meta with adlevel = DataOnly } } | FunApp (o, f, l) -> - let l = List.map ~f:(update_expr_ad_levels autodiffable_variables) l in - {pattern= FunApp (o, f, l); meta= {e.meta with adlevel= ad_level_sup l}} + let l = List.map ~f:(update_expr_ad_levels autodiffable_variables) l in + { pattern = FunApp (o, f, l); meta = { e.meta with adlevel = ad_level_sup l } } | TernaryIf (e1, e2, e3) -> - let e1 = update_expr_ad_levels autodiffable_variables e1 in - let e2 = update_expr_ad_levels autodiffable_variables e2 in - let e3 = update_expr_ad_levels autodiffable_variables e3 in - { pattern= TernaryIf (e1, e2, e3) - ; meta= {e.meta with adlevel= ad_level_sup [e1; e2; e3]} } + let e1 = update_expr_ad_levels autodiffable_variables e1 in + let e2 = update_expr_ad_levels autodiffable_variables e2 in + let e3 = update_expr_ad_levels autodiffable_variables e3 in + { pattern = TernaryIf (e1, e2, e3) + ; meta = { e.meta with adlevel = ad_level_sup [ e1; e2; e3 ] } + } | EAnd (e1, e2) -> - let e1 = update_expr_ad_levels autodiffable_variables e1 in - let e2 = update_expr_ad_levels autodiffable_variables e2 in - { pattern= EAnd (e1, e2) - ; meta= {e.meta with adlevel= ad_level_sup [e1; e2]} } + let e1 = update_expr_ad_levels autodiffable_variables e1 in + let e2 = update_expr_ad_levels autodiffable_variables e2 in + { pattern = EAnd (e1, e2); meta = { e.meta with adlevel = ad_level_sup [ e1; e2 ] } } | EOr (e1, e2) -> - let e1 = update_expr_ad_levels autodiffable_variables e1 in - let e2 = update_expr_ad_levels autodiffable_variables e2 in - { pattern= EOr (e1, e2) - ; meta= {e.meta with adlevel= ad_level_sup [e1; e2]} } + let e1 = update_expr_ad_levels autodiffable_variables e1 in + let e2 = update_expr_ad_levels autodiffable_variables e2 in + { pattern = EOr (e1, e2); meta = { e.meta with adlevel = ad_level_sup [ e1; e2 ] } } | Indexed (ixed, i_list) -> - let ixed = update_expr_ad_levels autodiffable_variables ixed in - let i_list = - List.map ~f:(update_idx_ad_levels autodiffable_variables) i_list - in - { pattern= Indexed (ixed, i_list) - ; meta= - { e.meta with - adlevel= ad_level_sup (e :: List.concat_map ~f:Index.bounds i_list) - } } + let ixed = update_expr_ad_levels autodiffable_variables ixed in + let i_list = List.map ~f:(update_idx_ad_levels autodiffable_variables) i_list in + { pattern = Indexed (ixed, i_list) + ; meta = + { e.meta with + adlevel = ad_level_sup (e :: List.concat_map ~f:Index.bounds i_list) + } + } and update_idx_ad_levels autodiffable_variables = Index.map (update_expr_ad_levels autodiffable_variables) +;; (** [cleanup_stmts statements] will do a few simple transformations like removing Skips, collapsing empty blocks and SLists, etc. *) @@ -432,34 +475,42 @@ let cleanup_empty_stmts stmts = let open Stmt.Fixed in let open Stmt.Fixed.Pattern in let cleanup_stmt s = - let ellide = {s with pattern= Skip} in + let ellide = { s with pattern = Skip } in match s.pattern with | Block [] | SList [] -> ellide - | For {body= {pattern= Skip; _}; _} -> ellide - | While (_, {pattern= Skip; _}) -> ellide - | Block [{pattern= Skip; _}] | SList [{pattern= Skip; _}] -> ellide + | For { body = { pattern = Skip; _ }; _ } -> ellide + | While (_, { pattern = Skip; _ }) -> ellide + | Block [ { pattern = Skip; _ } ] | SList [ { pattern = Skip; _ } ] -> ellide | _ -> s in - let is_decl = function {pattern= Decl _; _} -> true | _ -> false in + let is_decl = function + | { pattern = Decl _; _ } -> true + | _ -> false + in let flatten_block s = match s.pattern with - | SList ls | Block ls -> - if List.for_all ~f:(Fn.non is_decl) ls then ls else [s] - | _ -> [s] + | SList ls | Block ls -> if List.for_all ~f:(Fn.non is_decl) ls then ls else [ s ] + | _ -> [ s ] + in + let ellide_skip s = + match s.pattern with + | Skip -> [] + | _ -> [ s ] in - let ellide_skip s = match s.pattern with Skip -> [] | _ -> [s] in List.map stmts ~f:(rewrite_bottom_up ~f:Fn.id ~g:cleanup_stmt) |> List.concat_map ~f:flatten_block |> List.concat_map ~f:ellide_skip +;; let%expect_test "cleanup" = let open Expr.Helpers in let open Stmt.Fixed in let open Stmt.Fixed.Pattern in - let swrap pattern = {pattern; meta= Location_span.empty} in - let body = Block [Skip |> swrap] |> swrap in - let s = For {loopvar= "i"; lower= loop_bottom; upper= loop_bottom; body} in - let res = [s |> swrap] |> cleanup_empty_stmts in - [%sexp (res : Stmt.Located.t list)] |> print_s ; + let swrap pattern = { pattern; meta = Location_span.empty } in + let body = Block [ Skip |> swrap ] |> swrap in + let s = For { loopvar = "i"; lower = loop_bottom; upper = loop_bottom; body } in + let res = [ s |> swrap ] |> cleanup_empty_stmts in + [%sexp (res : Stmt.Located.t list)] |> print_s; [%expect {| () |}] +;; diff --git a/src/analysis_and_optimization/Mir_utils.mli b/src/analysis_and_optimization/Mir_utils.mli index d5bafaa765..71829f8465 100644 --- a/src/analysis_and_optimization/Mir_utils.mli +++ b/src/analysis_and_optimization/Mir_utils.mli @@ -6,85 +6,86 @@ val var_declarations : ('a, 'b) Stmt.Fixed.t -> string Set.Poly.t val num_expr_value : Expr.Typed.t -> (float * string) option type bound_values = - { lower: [`None | `Nonlit | `Lit of float] - ; upper: [`None | `Nonlit | `Lit of float] } + { lower : [ `None | `Nonlit | `Lit of float ] + ; upper : [ `None | `Nonlit | `Lit of float ] + } val trans_bounds_values : Expr.Typed.t Program.transformation -> bound_values val is_dist : string -> bool val chop_dist_name : string -> string Option.t val top_var_declarations : Stmt.Located.t -> string Set.Poly.t -val data_set : - ?exclude_transformed:bool +val data_set + : ?exclude_transformed:bool -> ?exclude_ints:bool -> Program.Typed.t -> string Set.Poly.t -val parameter_set : - ?include_transformed:bool +val parameter_set + : ?include_transformed:bool -> Program.Typed.t -> (string * Expr.Typed.t Program.transformation) Set.Poly.t -val parameter_names_set : - ?include_transformed:bool -> Program.Typed.t -> string Set.Poly.t +val parameter_names_set + : ?include_transformed:bool + -> Program.Typed.t + -> string Set.Poly.t -val fold_expr : - take_expr:('c -> Expr.Typed.Meta.t Expr.Fixed.t -> 'c) +val fold_expr + : take_expr:('c -> Expr.Typed.Meta.t Expr.Fixed.t -> 'c) -> init:'c -> Expr.Typed.t -> 'c -val fold_stmts : - take_expr:('c -> Expr.Typed.Meta.t Expr.Fixed.t -> 'c) +val fold_stmts + : take_expr:('c -> Expr.Typed.Meta.t Expr.Fixed.t -> 'c) -> take_stmt:('c -> Stmt.Located.t -> 'c) -> init:'c -> Stmt.Located.t List.t -> 'c -val map_rec_expr : - (Expr.Typed.t Expr.Fixed.Pattern.t -> Expr.Typed.t Expr.Fixed.Pattern.t) +val map_rec_expr + : (Expr.Typed.t Expr.Fixed.Pattern.t -> Expr.Typed.t Expr.Fixed.Pattern.t) -> Expr.Typed.t -> Expr.Typed.t -val map_rec_expr_state : - ( 's - -> Expr.Typed.t Expr.Fixed.Pattern.t - -> Expr.Typed.t Expr.Fixed.Pattern.t * 's) +val map_rec_expr_state + : ('s -> Expr.Typed.t Expr.Fixed.Pattern.t -> Expr.Typed.t Expr.Fixed.Pattern.t * 's) -> 's -> Expr.Typed.t -> Expr.Typed.t * 's -val map_rec_stmt_loc : - ( (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t +val map_rec_stmt_loc + : ((Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) -> Stmt.Located.t -> Stmt.Located.t -val top_down_map_rec_stmt_loc : - ( (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t +val top_down_map_rec_stmt_loc + : ((Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) -> Stmt.Located.t -> Stmt.Located.t -val map_rec_state_stmt_loc : - ( 's +val map_rec_state_stmt_loc + : ('s -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t * 's) -> 's -> Stmt.Located.t -> Stmt.Located.t * 's -val map_rec_stmt_loc_num : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t - -> ( int +val map_rec_stmt_loc_num + : (int, Stmt.Located.Non_recursive.t) Map.Poly.t + -> (int -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) -> Stmt.Located.Non_recursive.t -> Stmt.Located.t -val map_rec_state_stmt_loc_num : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t - -> ( int +val map_rec_state_stmt_loc_num + : (int, Stmt.Located.Non_recursive.t) Map.Poly.t + -> (int -> 's -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t * 's) @@ -92,123 +93,115 @@ val map_rec_state_stmt_loc_num : -> Stmt.Located.Non_recursive.t -> Stmt.Located.t * 's -val stmt_loc_of_stmt_loc_num : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t +val stmt_loc_of_stmt_loc_num + : (int, Stmt.Located.Non_recursive.t) Map.Poly.t -> Stmt.Located.Non_recursive.t -> Stmt.Located.t -val statement_stmt_loc_of_statement_stmt_loc_num : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t +val statement_stmt_loc_of_statement_stmt_loc_num + : (int, Stmt.Located.Non_recursive.t) Map.Poly.t -> (Expr.Typed.t, int) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t -val unnumbered_prog_of_numbered_prog : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t +val unnumbered_prog_of_numbered_prog + : (int, Stmt.Located.Non_recursive.t) Map.Poly.t -> ('a -> 'b) -> (Stmt.Located.Non_recursive.t, 'a) Program.t -> (Stmt.Located.t, 'b) Program.t -val fwd_traverse_statement : - ('e, 'a) Stmt.Fixed.Pattern.t - -> init:'f - -> f:('f -> 'a -> 'f * 'c) - -> 'f * ('e, 'c) Stmt.Fixed.Pattern.t (** A traversal that simultaneously accumulates a state (type 'f) and replaces the substatement values from ('a to 'c). Traversal is done in-order but ignores branching, e.g., and if's then block is followed by the else block rather than branching. *) +val fwd_traverse_statement + : ('e, 'a) Stmt.Fixed.Pattern.t + -> init:'f + -> f:('f -> 'a -> 'f * 'c) + -> 'f * ('e, 'c) Stmt.Fixed.Pattern.t -val vexpr_of_expr_exn : Expr.Typed.t -> vexpr (** Take a LHS expression from a general expression, throwing an exception if it can't be a LHS expression. *) +val vexpr_of_expr_exn : Expr.Typed.t -> vexpr -val expr_var_set : Expr.Typed.t -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t (** The set of variables in an expression, including inside an index. For use in RHS sets, not LHS assignment sets, except in a target term. *) +val expr_var_set : Expr.Typed.t -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t -val index_var_set : - Expr.Typed.t Index.t -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t (** The set of variables in an index. For use in RHS sets, not LHS assignment sets, except in a target term *) +val index_var_set : Expr.Typed.t Index.t -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t -val stmt_rhs : - (Expr.Typed.t, 's) Stmt.Fixed.Pattern.t -> Expr.Typed.t Set.Poly.t (** The set of variables that can affect the value or behavior of the expression, i.e. rhs. Using Set.Poly instead of ExprSet so that 'e can be polymorphic, it usually doesn't matter if there's duplication. *) +val stmt_rhs : (Expr.Typed.t, 's) Stmt.Fixed.Pattern.t -> Expr.Typed.t Set.Poly.t -val union_map : 'a Set.Poly.t -> f:('a -> 'b Set.Poly.t) -> 'b Set.Poly.t (** This is a helper function equivalent to List.concat_map but for Sets *) +val union_map : 'a Set.Poly.t -> f:('a -> 'b Set.Poly.t) -> 'b Set.Poly.t -val stmt_rhs_var_set : - (Expr.Typed.t, 's) Stmt.Fixed.Pattern.t - -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t (** The set of variables in an expression, including inside an index. For use in RHS sets, not LHS assignment sets, except in a target term. *) +val stmt_rhs_var_set + : (Expr.Typed.t, 's) Stmt.Fixed.Pattern.t + -> (vexpr * Expr.Typed.Meta.t) Set.Poly.t -val expr_assigned_var : Expr.Typed.t -> vexpr (** The variable being assigned to when the expression is the LHS *) +val expr_assigned_var : Expr.Typed.t -> vexpr -val summation_terms : Expr.Typed.t -> Expr.Typed.t list (** The list of terms in expression separated by a + *) +val summation_terms : Expr.Typed.t -> Expr.Typed.t list -val stmt_of_block : Stmt.Located.t list -> Stmt.Located.t (** Represent a list of statements as a single statement *) +val stmt_of_block : Stmt.Located.t list -> Stmt.Located.t -val subst_expr : - (string, Expr.Typed.t) Map.Poly.t -> Expr.Typed.t -> Expr.Typed.t (** Substitute variables in an expression according to the provided Map. *) +val subst_expr : (string, Expr.Typed.t) Map.Poly.t -> Expr.Typed.t -> Expr.Typed.t -val subst_stmt_base : - (string, Expr.Typed.t) Map.Poly.t +(** Substitute variables occurring at the top level in statements according to the provided Map. *) +val subst_stmt_base + : (string, Expr.Typed.t) Map.Poly.t -> (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t -(** Substitute variables occurring at the top level in statements according to the provided Map. *) -val subst_stmt : - (string, Expr.Typed.t) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t (** Substitute variables occurring anywhere in a statement according to the provided Map. *) +val subst_stmt : (string, Expr.Typed.t) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t -val name_subst_stmt : - (string, string) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t (** Substitute subexpressions occurring anywhere in a statement according to the provided Map. *) +val name_subst_stmt : (string, string) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t -val expr_subst_expr : - Expr.Typed.t Expr.Typed.Map.t -> Expr.Typed.t -> Expr.Typed.t (** Substitute subexpressions in an expression according to the provided Map, trying to match on larger subexpressions before smaller ones. *) +val expr_subst_expr : Expr.Typed.t Expr.Typed.Map.t -> Expr.Typed.t -> Expr.Typed.t -val expr_subst_stmt : - Expr.Typed.t Expr.Typed.Map.t -> Stmt.Located.t -> Stmt.Located.t (** Substitute subexpressions occurring anywhere in a statement according to the provided Map. *) +val expr_subst_stmt : Expr.Typed.t Expr.Typed.Map.t -> Stmt.Located.t -> Stmt.Located.t -val expr_subst_stmt_base : - Expr.Typed.t Expr.Typed.Map.t +(** Substitute subexpressions occurring at the top level in statements according to the provided Map. *) +val expr_subst_stmt_base + : Expr.Typed.t Expr.Typed.Map.t -> (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t -> (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t -(** Substitute subexpressions occurring at the top level in statements according to the provided Map. *) -val expr_depth : Expr.Typed.t -> int (** Calculate how deeply nested an expression is. *) +val expr_depth : Expr.Typed.t -> int -val update_expr_ad_levels : string Set.Poly.t -> Expr.Typed.t -> Expr.Typed.t (** Recompute all AD-levels in the metadata of an expression from the bottom up, making the variables in the first argument autodiffable *) +val update_expr_ad_levels : string Set.Poly.t -> Expr.Typed.t -> Expr.Typed.t -val cleanup_empty_stmts : - ('e, 's) Stmt.Fixed.t list -> ('e, 's) Stmt.Fixed.t list +val cleanup_empty_stmts : ('e, 's) Stmt.Fixed.t list -> ('e, 's) Stmt.Fixed.t list diff --git a/src/analysis_and_optimization/Monotone_framework.ml b/src/analysis_and_optimization/Monotone_framework.ml index 12d08d8fca..ec28a7f4ed 100644 --- a/src/analysis_and_optimization/Monotone_framework.ml +++ b/src/analysis_and_optimization/Monotone_framework.ml @@ -8,21 +8,26 @@ open Middle let preserve_stability = false (** Debugging tool to print out MFP sets **) -let print_mfp to_string (mfp : (int, 'a entry_exit) Map.Poly.t) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) : unit - = +let print_mfp + to_string + (mfp : (int, 'a entry_exit) Map.Poly.t) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + : unit + = let print_set s = - [%sexp (Set.Poly.map ~f:to_string s : string Set.Poly.t)] - |> Sexp.to_string_hum - in - let print_stmt s = - [%sexp (s : Stmt.Located.Non_recursive.t)] |> Sexp.to_string_hum + [%sexp (Set.Poly.map ~f:to_string s : string Set.Poly.t)] |> Sexp.to_string_hum in + let print_stmt s = [%sexp (s : Stmt.Located.Non_recursive.t)] |> Sexp.to_string_hum in Map.iteri mfp ~f:(fun ~key ~data -> print_endline - ( string_of_int key ^ ":\n " + (string_of_int key + ^ ":\n " ^ print_stmt (Map.Poly.find_exn flowgraph_to_mir key) - ^ ":\n " ^ print_set data.entry ^ " \t-> " ^ print_set data.exit ) ) + ^ ":\n " + ^ print_set data.entry + ^ " \t-> " + ^ print_set data.exit)) +;; (** Calculate the free (non-bound) variables in an expression *) let rec free_vars_expr (e : Expr.Typed.t) = @@ -30,13 +35,12 @@ let rec free_vars_expr (e : Expr.Typed.t) = | Var x -> Set.Poly.singleton x | Lit (_, _) -> Set.Poly.empty | FunApp (_, f, l) -> - Set.Poly.union_list (Set.Poly.singleton f :: List.map ~f:free_vars_expr l) + Set.Poly.union_list (Set.Poly.singleton f :: List.map ~f:free_vars_expr l) | TernaryIf (e1, e2, e3) -> - Set.Poly.union_list (List.map ~f:free_vars_expr [e1; e2; e3]) - | Indexed (e, l) -> - Set.Poly.union_list (free_vars_expr e :: List.map ~f:free_vars_idx l) + Set.Poly.union_list (List.map ~f:free_vars_expr [ e1; e2; e3 ]) + | Indexed (e, l) -> Set.Poly.union_list (free_vars_expr e :: List.map ~f:free_vars_idx l) | EAnd (e1, e2) | EOr (e1, e2) -> - Set.Poly.union_list (List.map ~f:free_vars_expr [e1; e2]) + Set.Poly.union_list (List.map ~f:free_vars_expr [ e1; e2 ]) (** Calculate the free (non-bound) variables in an index*) and free_vars_idx (i : Expr.Typed.t Index.t) = @@ -44,60 +48,64 @@ and free_vars_idx (i : Expr.Typed.t Index.t) = | All -> Set.Poly.empty | Single e | Upfrom e | MultiIndex e -> free_vars_expr e | Between (e1, e2) -> Set.Poly.union (free_vars_expr e1) (free_vars_expr e2) +;; (** Calculate the free (non-bound) variables in a statement *) -let rec free_vars_stmt - (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = +let rec free_vars_stmt (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = match s with - | Assignment ((_, _, []), e) | Return (Some e) | TargetPE e -> - free_vars_expr e + | Assignment ((_, _, []), e) | Return (Some e) | TargetPE e -> free_vars_expr e | Assignment ((_, _, l), e) -> - Set.Poly.union_list (free_vars_expr e :: List.map ~f:free_vars_idx l) + Set.Poly.union_list (free_vars_expr e :: List.map ~f:free_vars_idx l) | NRFunApp (_, f, l) -> - Set.Poly.union_list (Set.Poly.singleton f :: List.map ~f:free_vars_expr l) + Set.Poly.union_list (Set.Poly.singleton f :: List.map ~f:free_vars_expr l) | IfElse (e, b1, Some b2) -> - Set.Poly.union_list - [free_vars_expr e; free_vars_stmt b1.pattern; free_vars_stmt b2.pattern] + Set.Poly.union_list + [ free_vars_expr e; free_vars_stmt b1.pattern; free_vars_stmt b2.pattern ] | IfElse (e, b, None) | While (e, b) -> - Set.Poly.union (free_vars_expr e) (free_vars_stmt b.pattern) - | For {lower= e1; upper= e2; body= b; _} -> - Set.Poly.union_list - [free_vars_expr e1; free_vars_expr e2; free_vars_stmt b.pattern] + Set.Poly.union (free_vars_expr e) (free_vars_stmt b.pattern) + | For { lower = e1; upper = e2; body = b; _ } -> + Set.Poly.union_list [ free_vars_expr e1; free_vars_expr e2; free_vars_stmt b.pattern ] | Block l | SList l -> - Set.Poly.union_list (List.map ~f:(fun s -> free_vars_stmt s.pattern) l) + Set.Poly.union_list (List.map ~f:(fun s -> free_vars_stmt s.pattern) l) | Decl _ | Break | Continue | Return None | Skip -> Set.Poly.empty +;; (** A variation on free_vars_stmt, where we do not recursively count free variables in sub statements *) let top_free_vars_stmt (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) = + (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) + = match s with - | Assignment _ | Return _ | TargetPE _ | NRFunApp _ | Decl _ | Break - |Continue | Skip -> - free_vars_stmt - (statement_stmt_loc_of_statement_stmt_loc_num flowgraph_to_mir s) + | Assignment _ | Return _ | TargetPE _ | NRFunApp _ | Decl _ | Break | Continue | Skip + -> free_vars_stmt (statement_stmt_loc_of_statement_stmt_loc_num flowgraph_to_mir s) | While (e, _) | IfElse (e, _, _) -> free_vars_expr e - | For {lower= e1; upper= e2; _} -> - Set.Poly.union_list [free_vars_expr e1; free_vars_expr e2] + | For { lower = e1; upper = e2; _ } -> + Set.Poly.union_list [ free_vars_expr e1; free_vars_expr e2 ] | Block _ | SList _ -> Set.Poly.empty +;; (** Compute the inverse flowgraph of a Stan statement (for reverse analyses) *) -let inverse_flowgraph_of_stmt ?(flatten_loops = false) - ?(blocks_after_body = true) (stmt : Stmt.Located.t) : - (module FLOWGRAPH with type labels = int) - * (int, Stmt.Located.Non_recursive.t) Map.Poly.t = +let inverse_flowgraph_of_stmt + ?(flatten_loops = false) + ?(blocks_after_body = true) + (stmt : Stmt.Located.t) + : (module FLOWGRAPH with type labels = int) + * (int, Stmt.Located.Non_recursive.t) Map.Poly.t + = let flowgraph_to_mir = Dataflow_utils.build_statement_map - (fun Stmt.Fixed.({pattern; _}) -> pattern) - (fun Stmt.Fixed.({meta; _}) -> meta) + (fun Stmt.Fixed.{ pattern; _ } -> pattern) + (fun Stmt.Fixed.{ meta; _ } -> meta) stmt in let initials, successors = - Dataflow_utils.build_predecessor_graph ~flatten_loops ~blocks_after_body + Dataflow_utils.build_predecessor_graph + ~flatten_loops + ~blocks_after_body flowgraph_to_mir in - ( ( module struct + ( (module struct type labels = int type t = labels @@ -106,17 +114,17 @@ let inverse_flowgraph_of_stmt ?(flatten_loops = false) let sexp_of_t = Int.sexp_of_t let initials = initials let successors = successors - end - : FLOWGRAPH - with type labels = int ) + end : FLOWGRAPH + with type labels = int) , Map.Poly.map - ~f:(fun (pattern, meta) -> Stmt.Located.Non_recursive.{pattern; meta}) + ~f:(fun (pattern, meta) -> Stmt.Located.Non_recursive.{ pattern; meta }) flowgraph_to_mir ) +;; (** Reverse flowgraphs to be used for reverse analyses. Observe that this respects the invariants listed for a FLOWGRAPH *) let reverse (type l) (module F : FLOWGRAPH with type labels = l) = - ( module struct + (module struct type labels = F.labels type t = labels @@ -126,172 +134,199 @@ let reverse (type l) (module F : FLOWGRAPH with type labels = l) = let initials = Set.of_map_keys (Map.filter F.successors ~f:Set.is_empty) let successors = - Map.fold F.successors + Map.fold + F.successors ~init:(Map.map F.successors ~f:(fun _ -> Set.Poly.empty)) ~f:(fun ~key:old_pred ~data:old_succs accum -> Set.fold old_succs ~init:accum ~f:(fun accum old_succ -> - Map.set accum ~key:old_succ - ~data:(Set.add (Map.find_exn accum old_succ) old_pred) ) ) - end - : FLOWGRAPH - with type labels = l ) + Map.set + accum + ~key:old_succ + ~data:(Set.add (Map.find_exn accum old_succ) old_pred))) + ;; + end : FLOWGRAPH + with type labels = l) +;; (** Compute the forward flowgraph of a Stan statement (for forward analyses) *) -let forward_flowgraph_of_stmt ?(flatten_loops = false) - ?(blocks_after_body = true) stmt = - let inv_flowgraph = - inverse_flowgraph_of_stmt ~flatten_loops ~blocks_after_body stmt - in - (reverse (fst inv_flowgraph), snd inv_flowgraph) +let forward_flowgraph_of_stmt ?(flatten_loops = false) ?(blocks_after_body = true) stmt = + let inv_flowgraph = inverse_flowgraph_of_stmt ~flatten_loops ~blocks_after_body stmt in + reverse (fst inv_flowgraph), snd inv_flowgraph +;; (** The lattice of sets of some values, with the inclusion order, set union and the empty set *) let powerset_lattice (type v) (module S : INITIALTYPE with type vals = v) = - ( module struct + (module struct type properties = S.vals Set.Poly.t let bottom = Set.Poly.empty let lub s1 s2 = Set.Poly.union s1 s2 let leq s1 s2 = Set.Poly.is_subset s1 ~of_:s2 let initial = S.initial - end - : LATTICE - with type properties = v Set.Poly.t ) + end : LATTICE + with type properties = v Set.Poly.t) +;; (** The lattice of subsets of some set, with the inverse inclusion order, set intersection and the total set *) -let dual_powerset_lattice (type v) - (module S : INITIALTOTALTYPE with type vals = v) = - ( module struct +let dual_powerset_lattice (type v) (module S : INITIALTOTALTYPE with type vals = v) = + (module struct type properties = S.vals Set.Poly.t let bottom = S.total let lub s1 s2 = Set.Poly.inter s1 s2 let leq s1 s2 = Set.Poly.is_subset s2 ~of_:s1 let initial = S.initial - end - : LATTICE - with type properties = v Set.Poly.t ) + end : LATTICE + with type properties = v Set.Poly.t) +;; let powerset_lattice_expressions (initial : Expr.Typed.Set.t) = - ( module struct + (module struct type properties = Expr.Typed.Set.t let bottom = Expr.Typed.Set.empty let lub s1 s2 = Expr.Typed.Set.union s1 s2 let leq s1 s2 = Expr.Typed.Set.is_subset s1 ~of_:s2 let initial = initial - end - : LATTICE - with type properties = Expr.Typed.Set.t ) - -let dual_powerset_lattice_expressions (initial : Expr.Typed.Set.t) - (total : Expr.Typed.Set.t) = - ( module struct + end : LATTICE + with type properties = Expr.Typed.Set.t) +;; + +let dual_powerset_lattice_expressions + (initial : Expr.Typed.Set.t) + (total : Expr.Typed.Set.t) + = + (module struct type properties = Expr.Typed.Set.t let bottom = total let lub s1 s2 = Expr.Typed.Set.inter s1 s2 let leq s1 s2 = Expr.Typed.Set.is_subset s2 ~of_:s1 let initial = initial - end - : LATTICE - with type properties = Expr.Typed.Set.t ) + end : LATTICE + with type properties = Expr.Typed.Set.t) +;; (** Add a fresh bottom element to a lattice (possibly without bottom) *) let new_bot (type p) (module L : LATTICE_NO_BOT with type properties = p) = - ( module struct + (module struct type properties = L.properties option let bottom = None let lub = function - | Some s1 -> ( - function Some s2 -> Some (L.lub s1 s2) | None -> Some s1 ) + | Some s1 -> + (function + | Some s2 -> Some (L.lub s1 s2) + | None -> Some s1) | None -> fun x -> x + ;; let leq = function - | Some s1 -> ( function Some s2 -> L.leq s1 s2 | None -> false ) + | Some s1 -> + (function + | Some s2 -> L.leq s1 s2 + | None -> false) | None -> fun _ -> true + ;; let initial = Some L.initial - end - : LATTICE - with type properties = p option ) + end : LATTICE + with type properties = p option) +;; (** The lattice (without bottom) of partial functions, ordered under inverse graph inclusion, with intersection *) -let dual_partial_function_lattice (type dv cv) +let dual_partial_function_lattice + (type dv cv) (module Dom : TOTALTYPE with type vals = dv) - (module Codom : TYPE with type vals = cv) = - ( module struct + (module Codom : TYPE with type vals = cv) + = + (module struct type properties = (Dom.vals, Codom.vals) Map.Poly.t (* intersection *) let lub s1 s2 = let f ~key ~data = Map.find s2 key = Some data in Map.filteri ~f s1 + ;; let leq s1 s2 = Set.for_all Dom.total ~f:(fun k -> - match (Map.find s1 k, Map.find s2 k) with + match Map.find s1 k, Map.find s2 k with | Some x, Some y -> x = y | Some _, None | None, None -> true - | None, Some _ -> false ) + | None, Some _ -> false) + ;; let initial = Map.Poly.empty - end - : LATTICE_NO_BOT - with type properties = (dv, cv) Map.Poly.t ) + end : LATTICE_NO_BOT + with type properties = (dv, cv) Map.Poly.t) +;; (* The lattice of partial functions, where we add a fresh bottom element, to represent an inconsistent combination of functions *) -let dual_partial_function_lattice_with_bot (type dv cv) +let dual_partial_function_lattice_with_bot + (type dv cv) (module Dom : TOTALTYPE with type vals = dv) - (module Codom : TYPE with type vals = cv) = + (module Codom : TYPE with type vals = cv) + = new_bot (dual_partial_function_lattice (module Dom) (module Codom)) +;; (* A dual powerset lattice, where we set the initial set to be empty *) -let dual_powerset_lattice_empty_initial (type v) - (module T : TOTALTYPE with type vals = v) = +let dual_powerset_lattice_empty_initial (type v) (module T : TOTALTYPE with type vals = v) + = dual_powerset_lattice - ( module struct + (module struct type vals = T.vals let initial = Set.Poly.empty let total = T.total - end ) + end) +;; (* A powerset lattice, where we set the initial set to be empty *) -let powerset_lattice_empty_initial (type v) - (module T : TYPE with type vals = v) = +let powerset_lattice_empty_initial (type v) (module T : TYPE with type vals = v) = powerset_lattice - (module struct type vals = T.vals + (module struct + type vals = T.vals - let initial = Set.Poly.empty end) + let initial = Set.Poly.empty + end) +;; (* The specific powerset lattice we use for reaching definitions analysis *) -let reaching_definitions_lattice (type v l) +let reaching_definitions_lattice + (type v l) (module Variables : INITIALTYPE with type vals = v) - (module Labels : TYPE with type vals = l) = + (module Labels : TYPE with type vals = l) + = powerset_lattice - ( module struct + (module struct type vals = Variables.vals * Labels.vals option - let initial = Set.Poly.map ~f:(fun x -> (x, None)) Variables.initial - end ) + let initial = Set.Poly.map ~f:(fun x -> x, None) Variables.initial + end) +;; (* Autodiff-level lattice *) let autodiff_level_lattice autodiff_variables = powerset_lattice - (module struct type vals = string + (module struct + type vals = string - let initial = autodiff_variables end) + let initial = autodiff_variables + end) +;; (* The transfer function for a constant propagation analysis *) let constant_propagation_transfer - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = (string, Expr.Typed.t) Map.Poly.t option @@ -299,47 +334,48 @@ let constant_propagation_transfer match p with | None -> None | Some m -> - let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in - Some - ( match mir_node with - (* TODO: we are currently only propagating constants for scalars. + let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in + Some + (match mir_node with + (* TODO: we are currently only propagating constants for scalars. We could do the same for matrix and array expressions if we wanted. *) - | Assignment ((s, t, []), e) -> ( - match Partial_evaluator.eval_expr (subst_expr m e) with - | {pattern= Lit (_, _); _} as e' - when not (preserve_stability && UnsizedType.is_autodiffable t) - -> - Map.set m ~key:s ~data:e' - | _ -> Map.remove m s ) - | Decl {decl_id= s; _} | Assignment ((s, _, _ :: _), _) -> - Map.remove m s - | TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |For _ | Block _ | SList _ -> - m ) - end - : TRANSFER_FUNCTION + | Assignment ((s, t, []), e) -> + (match Partial_evaluator.eval_expr (subst_expr m e) with + | { pattern = Lit (_, _); _ } as e' + when not (preserve_stability && UnsizedType.is_autodiffable t) -> + Map.set m ~key:s ~data:e' + | _ -> Map.remove m s) + | Decl { decl_id = s; _ } | Assignment ((s, _, _ :: _), _) -> Map.remove m s + | TargetPE _ + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | For _ | Block _ | SList _ -> m) + ;; + end : TRANSFER_FUNCTION with type labels = int - and type properties = (string, Expr.Typed.t) Map.Poly.t option ) + and type properties = (string, Expr.Typed.t) Map.Poly.t option) +;; let label_top_decls (flowgraph_to_mir : (int, Middle.Stmt.Located.Non_recursive.t) Map.Poly.t) - label : string Set.Poly.t = + label + : string Set.Poly.t + = let stmt = Map.Poly.find_exn flowgraph_to_mir label in match stmt.pattern with - | Decl {decl_id= s; _} -> Set.Poly.singleton s + | Decl { decl_id = s; _ } -> Set.Poly.singleton s | _ -> Set.Poly.empty +;; (** The transfer function for an expression propagation analysis, AKA forward substitution (see page 396 of Muchnick) *) let expression_propagation_transfer (can_side_effect_expr : Middle.Expr.Typed.t -> bool) (flowgraph_to_mir : (int, Middle.Stmt.Located.Non_recursive.t) Map.Poly.t) - = - ( module struct + = + (module struct type labels = int type properties = (string, Expr.Typed.t) Map.Poly.t option @@ -347,47 +383,46 @@ let expression_propagation_transfer match p with | None -> None | Some m -> - let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in - let kill_var m v = - Map.filteri m ~f:(fun ~key ~data -> - not (key = v || Set.Poly.mem (free_vars_expr data) v) ) - in - Some - ( match mir_node with - (* TODO: we are currently only propagating constants for scalars. + let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in + let kill_var m v = + Map.filteri m ~f:(fun ~key ~data -> + not (key = v || Set.Poly.mem (free_vars_expr data) v)) + in + Some + (match mir_node with + (* TODO: we are currently only propagating constants for scalars. We could do the same for matrix and array expressions if we wanted. *) - | Middle.Stmt.Fixed.Pattern.Assignment ((s, t, []), e) -> - let m' = kill_var m s in - if - can_side_effect_expr e - || Set.Poly.mem (free_vars_expr e) s - || (preserve_stability && UnsizedType.is_autodiffable t) - then m' - else Map.set m ~key:s ~data:(subst_expr m e) - | Decl {decl_id= s; _} | Assignment ((s, _, _ :: _), _) -> - kill_var m s - | Block b -> - let kills = - Set.Poly.union_list - (List.map ~f:(label_top_decls flowgraph_to_mir) b) - in - Set.Poly.fold kills ~init:m ~f:kill_var - | TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |For _ | SList _ -> - m ) - end - : TRANSFER_FUNCTION + | Middle.Stmt.Fixed.Pattern.Assignment ((s, t, []), e) -> + let m' = kill_var m s in + if can_side_effect_expr e + || Set.Poly.mem (free_vars_expr e) s + || (preserve_stability && UnsizedType.is_autodiffable t) + then m' + else Map.set m ~key:s ~data:(subst_expr m e) + | Decl { decl_id = s; _ } | Assignment ((s, _, _ :: _), _) -> kill_var m s + | Block b -> + let kills = + Set.Poly.union_list (List.map ~f:(label_top_decls flowgraph_to_mir) b) + in + Set.Poly.fold kills ~init:m ~f:kill_var + | TargetPE _ + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | For _ | SList _ -> m) + ;; + end : TRANSFER_FUNCTION with type labels = int - and type properties = (string, Expr.Typed.t) Map.Poly.t option ) + and type properties = (string, Expr.Typed.t) Map.Poly.t option) +;; (** The transfer function for a copy propagation analysis *) -let copy_propagation_transfer (globals : string Set.Poly.t) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct +let copy_propagation_transfer + (globals : string Set.Poly.t) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = (string, Expr.Typed.t) Map.Poly.t option @@ -395,35 +430,35 @@ let copy_propagation_transfer (globals : string Set.Poly.t) match p with | None -> None | Some m -> - let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in - let kill_var m v = - Map.filteri m ~f:(fun ~key ~(data : Expr.Typed.t) -> - not (key = v || data.pattern = Var v) ) - in - Some - ( match mir_node with - | Assignment ((s, _, []), {pattern= Var t; meta}) -> - let m' = kill_var m s in - if Set.Poly.mem globals s then m' - else Map.set m' ~key:s ~data:Expr.Fixed.{pattern= Var t; meta} - | Decl {decl_id= s; _} | Assignment ((s, _, _), _) -> kill_var m s - | Block b -> - let kills = - Set.Poly.union_list - (List.map ~f:(label_top_decls flowgraph_to_mir) b) - in - Set.Poly.fold kills ~init:m ~f:kill_var - | TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |For _ | SList _ -> - m ) - end - : TRANSFER_FUNCTION + let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in + let kill_var m v = + Map.filteri m ~f:(fun ~key ~(data : Expr.Typed.t) -> + not (key = v || data.pattern = Var v)) + in + Some + (match mir_node with + | Assignment ((s, _, []), { pattern = Var t; meta }) -> + let m' = kill_var m s in + if Set.Poly.mem globals s + then m' + else Map.set m' ~key:s ~data:Expr.Fixed.{ pattern = Var t; meta } + | Decl { decl_id = s; _ } | Assignment ((s, _, _), _) -> kill_var m s + | Block b -> + let kills = + Set.Poly.union_list (List.map ~f:(label_top_decls flowgraph_to_mir) b) + in + Set.Poly.fold kills ~init:m ~f:kill_var + | TargetPE _ + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | For _ | SList _ -> m) + ;; + end : TRANSFER_FUNCTION with type labels = int - and type properties = (string, Expr.Typed.t) Map.Poly.t option ) + and type properties = (string, Expr.Typed.t) Map.Poly.t option) +;; (** A helper function for building transfer functions from gen and kill sets *) let transfer_gen_kill p gen kill = Set.union gen (Set.diff p kill) @@ -435,68 +470,66 @@ let assigned_vars_stmt (s : (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t) = match s with | Assignment ((x, _, _), _) -> Set.Poly.singleton x | TargetPE _ -> Set.Poly.singleton "target" - | NRFunApp (_, s, _) when String.suffix s 3 = "_lp" -> - Set.Poly.singleton "target" - | For {loopvar= x; _} -> Set.Poly.singleton x - | Decl {decl_id= _; _} - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |Block _ | SList _ -> - Set.Poly.empty + | NRFunApp (_, s, _) when String.suffix s 3 = "_lp" -> Set.Poly.singleton "target" + | For { loopvar = x; _ } -> Set.Poly.singleton x + | Decl { decl_id = _; _ } + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | Block _ | SList _ -> Set.Poly.empty +;; (** Calculate the set of variables that a statement can declare *) let declared_vars_stmt (s : (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t) = match s with - | Decl {decl_id= x; _} -> Set.Poly.singleton x + | Decl { decl_id = x; _ } -> Set.Poly.singleton x | _ -> Set.Poly.empty +;; (** Calculate the set of variables that a statement can assign to or declare *) -let assigned_or_declared_vars_stmt - (s : (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t) = +let assigned_or_declared_vars_stmt (s : (Expr.Typed.t, 'a) Stmt.Fixed.Pattern.t) = Set.Poly.union (assigned_vars_stmt s) (declared_vars_stmt s) +;; (** The transfer function for a reaching definitions analysis *) let reaching_definitions_transfer - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = (string * labels option) Set.Poly.t let transfer_function l p = let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in let gen = - Set.Poly.map - ~f:(fun x -> (x, Some l)) - (assigned_or_declared_vars_stmt mir_node) + Set.Poly.map ~f:(fun x -> x, Some l) (assigned_or_declared_vars_stmt mir_node) in let kill = match mir_node with - | Decl {decl_id= x; _} - |Assignment ((x, _, []), _) - |For {loopvar= x; _} -> - Set.filter p ~f:(fun (y, _) -> y = x) + | Decl { decl_id = x; _ } | Assignment ((x, _, []), _) | For { loopvar = x; _ } -> + Set.filter p ~f:(fun (y, _) -> y = x) | TargetPE _ -> Set.filter p ~f:(fun (y, _) -> y = "target") | NRFunApp (_, s, _) when String.suffix s 3 = "_lp" -> - Set.filter p ~f:(fun (y, _) -> y = "target") + Set.filter p ~f:(fun (y, _) -> y = "target") | NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |Block _ | SList _ | Assignment _ -> - Set.Poly.empty + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | Block _ | SList _ | Assignment _ -> Set.Poly.empty in transfer_gen_kill p gen kill - end - : TRANSFER_FUNCTION + ;; + end : TRANSFER_FUNCTION with type labels = int - and type properties = (string * int option) Set.Poly.t ) + and type properties = (string * int option) Set.Poly.t) +;; (** The transfer function for an initialized variables analysis *) let initialized_vars_transfer - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = string Set.Poly.t @@ -504,14 +537,18 @@ let initialized_vars_transfer let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in let gen = assigned_vars_stmt mir_node in transfer_gen_kill p gen Set.Poly.empty - end - : TRANSFER_FUNCTION - with type labels = int and type properties = string Set.Poly.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = string Set.Poly.t) +;; (** The transfer function for a live variables analysis *) -let live_variables_transfer (never_kill : string Set.Poly.t) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct +let live_variables_transfer + (never_kill : string Set.Poly.t) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = string Set.Poly.t @@ -520,132 +557,138 @@ let live_variables_transfer (never_kill : string Set.Poly.t) let gen = top_free_vars_stmt flowgraph_to_mir mir_node in let kill = match mir_node with - | Assignment ((x, _, []), _) | Decl {decl_id= x; _} -> - Set.Poly.singleton x + | Assignment ((x, _, []), _) | Decl { decl_id = x; _ } -> Set.Poly.singleton x | TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip - |IfElse (_, _, _) - |While (_, _) - |For _ | Block _ | SList _ - |Assignment ((_, _, _ :: _), _) -> - Set.Poly.empty + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip + | IfElse (_, _, _) + | While (_, _) + | For _ | Block _ | SList _ + | Assignment ((_, _, _ :: _), _) -> Set.Poly.empty in transfer_gen_kill p gen (Set.Poly.diff kill never_kill) - end - : TRANSFER_FUNCTION - with type labels = int and type properties = string Set.Poly.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = string Set.Poly.t) +;; (** Calculate the set of sub-expressions of an expression *) let rec used_subexpressions_expr (e : Expr.Typed.t) = Expr.Typed.Set.union (Expr.Typed.Set.singleton e) - ( match e.pattern with + (match e.pattern with | Var _ | Lit (_, _) -> Expr.Typed.Set.empty | FunApp (_, _, l) -> - Expr.Typed.Set.union_list (List.map ~f:used_subexpressions_expr l) + Expr.Typed.Set.union_list (List.map ~f:used_subexpressions_expr l) | TernaryIf (e1, e2, e3) -> - Expr.Typed.Set.union_list - [ used_subexpressions_expr e1 - ; used_subexpressions_expr e2 - ; used_subexpressions_expr e3 ] + Expr.Typed.Set.union_list + [ used_subexpressions_expr e1 + ; used_subexpressions_expr e2 + ; used_subexpressions_expr e3 + ] | Indexed (e, l) -> - Expr.Typed.Set.union_list - ( used_subexpressions_expr e - :: List.map ~f:(used_expressions_idx_help used_subexpressions_expr) l - ) + Expr.Typed.Set.union_list + (used_subexpressions_expr e + :: List.map ~f:(used_expressions_idx_help used_subexpressions_expr) l) | EAnd (e1, e2) | EOr (e1, e2) -> - Expr.Typed.Set.union_list - [used_subexpressions_expr e1; used_subexpressions_expr e2] ) + Expr.Typed.Set.union_list + [ used_subexpressions_expr e1; used_subexpressions_expr e2 ]) and used_expressions_idx_help f (i : Expr.Typed.t Index.t) = match i with | All -> Expr.Typed.Set.empty | Single e | Upfrom e | MultiIndex e -> f e | Between (e1, e2) -> Expr.Typed.Set.union (f e1) (f e2) +;; (** Calculate the set of expressions of an expression *) let used_expressions_expr e = Expr.Typed.Set.singleton e -let rec used_expressions_stmt_help f - (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = +let rec used_expressions_stmt_help + f + (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) + = match s with | Assignment ((_, _, []), e) | TargetPE e | Return (Some e) -> f e | Assignment ((_, _, l), e) -> - Expr.Typed.Set.union (f e) - (Expr.Typed.Set.union_list - (List.map ~f:(used_expressions_idx_help f) l)) + Expr.Typed.Set.union + (f e) + (Expr.Typed.Set.union_list (List.map ~f:(used_expressions_idx_help f) l)) | IfElse (e, b1, Some b2) -> - Expr.Typed.Set.union_list - [ f e - ; used_expressions_stmt_help f b1.pattern - ; used_expressions_stmt_help f b2.pattern ] + Expr.Typed.Set.union_list + [ f e + ; used_expressions_stmt_help f b1.pattern + ; used_expressions_stmt_help f b2.pattern + ] | NRFunApp (_, _, l) -> Expr.Typed.Set.union_list (List.map ~f l) | Decl _ | Return None | Break | Continue | Skip -> Expr.Typed.Set.empty | IfElse (e, b, None) | While (e, b) -> - Expr.Typed.Set.union (f e) (used_expressions_stmt_help f b.pattern) - | For {lower= e1; upper= e2; body= b; loopvar= s} -> - Expr.Typed.Set.union_list - [ f e1; f e2 - ; used_expressions_stmt_help f b.pattern - ; Expr.Typed.Set.singleton - { pattern= Var s - ; meta= - Expr.Typed.Meta. - {type_= UInt; adlevel= DataOnly; loc= Location_span.empty} } - ] + Expr.Typed.Set.union (f e) (used_expressions_stmt_help f b.pattern) + | For { lower = e1; upper = e2; body = b; loopvar = s } -> + Expr.Typed.Set.union_list + [ f e1 + ; f e2 + ; used_expressions_stmt_help f b.pattern + ; Expr.Typed.Set.singleton + { pattern = Var s + ; meta = + Expr.Typed.Meta. + { type_ = UInt; adlevel = DataOnly; loc = Location_span.empty } + } + ] | Block l | SList l -> - Expr.Typed.Set.union_list - (List.map ~f:(fun s -> used_expressions_stmt_help f s.pattern) l) + Expr.Typed.Set.union_list + (List.map ~f:(fun s -> used_expressions_stmt_help f s.pattern) l) +;; (** Calculate the set of sub-expressions in a statement *) -let used_subexpressions_stmt = - used_expressions_stmt_help used_subexpressions_expr +let used_subexpressions_stmt = used_expressions_stmt_help used_subexpressions_expr (** Calculate the set of expressions in a statement *) let used_expressions_stmt = used_expressions_stmt_help used_expressions_expr -let top_used_expressions_stmt_help f - (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) = +let top_used_expressions_stmt_help f (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) = match s with | Assignment ((_, _, []), e) | TargetPE e | Return (Some e) -> f e | Assignment ((_, _, l), e) -> - Expr.Typed.Set.union (f e) - (Expr.Typed.Set.union_list - (List.map ~f:(used_expressions_idx_help f) l)) + Expr.Typed.Set.union + (f e) + (Expr.Typed.Set.union_list (List.map ~f:(used_expressions_idx_help f) l)) | While (e, _) | IfElse (e, _, _) -> f e | NRFunApp (_, _, l) -> Expr.Typed.Set.union_list (List.map ~f l) | Block _ | SList _ | Decl _ | Return None | Break | Continue | Skip -> - Expr.Typed.Set.empty - | For {lower= e1; upper= e2; _} -> Expr.Typed.Set.union_list [f e1; f e2] + Expr.Typed.Set.empty + | For { lower = e1; upper = e2; _ } -> Expr.Typed.Set.union_list [ f e1; f e2 ] +;; (** Calculate the set of sub-expressions at the top level in a statement *) -let top_used_subexpressions_stmt = - top_used_expressions_stmt_help used_subexpressions_expr +let top_used_subexpressions_stmt = top_used_expressions_stmt_help used_subexpressions_expr (** Calculate the set of expressions at the top level in a statement *) -let top_used_expressions_stmt = - top_used_expressions_stmt_help used_expressions_expr +let top_used_expressions_stmt = top_used_expressions_stmt_help used_expressions_expr (** Calculate the subset (of p) of expressions that will need to be recomputed as a consequence of evaluating the statement s (because of writes to variables performed by s) *) -let killed_expressions_stmt (p : Expr.Typed.Set.t) - (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) = +let killed_expressions_stmt + (p : Expr.Typed.Set.t) + (s : (Expr.Typed.t, int) Stmt.Fixed.Pattern.t) + = Expr.Typed.Set.filter p ~f:(fun e -> let free_vars = free_vars_expr e in (* Note: a simple test for membership would be more efficient here, but it would require us to duplicate some code. *) let assigned_vars = assigned_or_declared_vars_stmt s in - not (Set.Poly.is_empty (Set.Poly.inter free_vars assigned_vars)) ) + not (Set.Poly.is_empty (Set.Poly.inter free_vars assigned_vars))) +;; (** Calculate the set of subexpressions that needs to be computed at each node in the flowgraph *) let used (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - Map.Poly.fold flowgraph_to_mir ~init:Map.Poly.empty - ~f:(fun ~key ~data accum -> - Map.Poly.set accum ~key ~data:(top_used_subexpressions_stmt data.pattern) - ) + Map.Poly.fold flowgraph_to_mir ~init:Map.Poly.empty ~f:(fun ~key ~data accum -> + Map.Poly.set accum ~key ~data:(top_used_subexpressions_stmt data.pattern)) +;; (* TODO: figure out whether we will also want to reuse the computation of killed *) @@ -653,8 +696,9 @@ let used (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = code motion) *) let anticipated_expressions_transfer (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (used : (int, Expr.Typed.Set.t) Map.Poly.t) = - ( module struct + (used : (int, Expr.Typed.Set.t) Map.Poly.t) + = + (module struct type labels = int type properties = Expr.Typed.Set.t @@ -663,9 +707,11 @@ let anticipated_expressions_transfer let gen = Map.Poly.find_exn used l in let kill = killed_expressions_stmt p mir_node in transfer_gen_kill p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = Expr.Typed.Set.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = Expr.Typed.Set.t) +;; (** A helper function for defining transfer functions in terms of gen and kill sets in an alternative way, that is used in some of the subanalyses of lazy code motion *) @@ -680,39 +726,42 @@ let transfer_gen_kill_alt p gen kill = Set.diff (Set.union p gen) kill (** An available expressions analysis, to be used in lazy code motion *) let available_expressions_transfer (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (anticipated_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) = - ( module struct + (anticipated_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) + = + (module struct type labels = int type properties = Expr.Typed.Set.t let transfer_function l p = let mir_node = (Map.find_exn flowgraph_to_mir l).pattern in let gen = (Map.find_exn anticipated_expressions l).exit in - let kill = - killed_expressions_stmt (Expr.Typed.Set.union p gen) mir_node - in + let kill = killed_expressions_stmt (Expr.Typed.Set.union p gen) mir_node in transfer_gen_kill_alt p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = Expr.Typed.Set.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = Expr.Typed.Set.t) +;; (** Calculates the set of expressions that can be calculated for the first time at each node in the flow graph *) let earliest (anticipated_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) - (available_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) = - Map.fold anticipated_expressions ~init:Map.Poly.empty - ~f:(fun ~key ~data accum -> - Map.set accum ~key - ~data: - (Set.diff data.exit (Map.find_exn available_expressions key).entry) - ) + (available_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) + = + Map.fold anticipated_expressions ~init:Map.Poly.empty ~f:(fun ~key ~data accum -> + Map.set + accum + ~key + ~data:(Set.diff data.exit (Map.find_exn available_expressions key).entry)) +;; (** The transfer function for a postponable expressions analysis (as a part of lazy code motion) *) let postponable_expressions_transfer (earliest : (int, Expr.Typed.Set.t) Map.Poly.t) - (used : (int, Expr.Typed.Set.t) Map.Poly.t) = - ( module struct + (used : (int, Expr.Typed.Set.t) Map.Poly.t) + = + (module struct type labels = int type properties = Expr.Typed.Set.t @@ -720,15 +769,19 @@ let postponable_expressions_transfer let gen = Map.find_exn earliest l in let kill = Map.find_exn used l in transfer_gen_kill_alt p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = Expr.Typed.Set.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = Expr.Typed.Set.t) +;; (** Calculates the set of expressions that can be computed at the latest at each node *) -let latest (successors : (int, int Set.Poly.t) Map.Poly.t) +let latest + (successors : (int, int Set.Poly.t) Map.Poly.t) (earliest : (int, Expr.Typed.Set.t) Map.Poly.t) (postponable_expressions : (int, Expr.Typed.Set.t entry_exit) Map.Poly.t) - (used : (int, Expr.Typed.Set.t) Map.Poly.t) = + (used : (int, Expr.Typed.Set.t) Map.Poly.t) + = let earliest_or_postponable key = Expr.Typed.Set.union (Map.Poly.find_exn earliest key) @@ -738,16 +791,18 @@ let latest (successors : (int, int Set.Poly.t) Map.Poly.t) Set.filter (earliest_or_postponable key) ~f:(fun e -> Set.mem (Map.Poly.find_exn used key) e || Set.Poly.exists (Map.Poly.find_exn successors key) ~f:(fun s -> - not (Set.mem (earliest_or_postponable s) e) ) ) + not (Set.mem (earliest_or_postponable s) e))) in Map.fold successors ~init:Map.Poly.empty ~f:(fun ~key ~data:_ accum -> - Map.set accum ~key ~data:(latest key) ) + Map.set accum ~key ~data:(latest key)) +;; (** The transfer function for a used-not-latest expressions analysis, as a part of lazy code motion *) let used_not_latest_expressions_transfer (used : (int, Expr.Typed.Set.t) Map.Poly.t) - (latest : (int, Expr.Typed.Set.t) Map.Poly.t) = - ( module struct + (latest : (int, Expr.Typed.Set.t) Map.Poly.t) + = + (module struct type labels = int type properties = Expr.Typed.Set.t @@ -755,14 +810,17 @@ let used_not_latest_expressions_transfer let gen = Map.find_exn used l in let kill = Map.find_exn latest l in transfer_gen_kill_alt p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = Expr.Typed.Set.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = Expr.Typed.Set.t) +;; (** The transfer function for the first forward analysis part of determining optimal ad-levels for variables *) let autodiff_level_fwd1_transfer - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = - ( module struct + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = + (module struct type labels = int type properties = string Set.Poly.t @@ -771,27 +829,28 @@ let autodiff_level_fwd1_transfer let gen = match mir_node with | Assignment ((x, _, _), e) - when Expr.Typed.adlevel_of (update_expr_ad_levels p e) = AutoDiffable - -> - Set.Poly.singleton x + when Expr.Typed.adlevel_of (update_expr_ad_levels p e) = AutoDiffable -> + Set.Poly.singleton x | _ -> Set.Poly.empty in let kill = match mir_node with - | Decl {decl_id; decl_adtype= DataOnly; _} -> - Set.Poly.singleton decl_id + | Decl { decl_id; decl_adtype = DataOnly; _ } -> Set.Poly.singleton decl_id | _ -> Set.Poly.empty in transfer_gen_kill p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = string Set.Poly.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = string Set.Poly.t) +;; (** The transfer function for the reverse analysis part of determining optimal ad-levels for variables *) let autodiff_level_rev_transfer (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (fwd_ad_levels : (int, string Set.Poly.t) Map.Poly.t) = - ( module struct + (fwd_ad_levels : (int, string Set.Poly.t) Map.Poly.t) + = + (module struct type labels = int type properties = string Set.Poly.t @@ -800,21 +859,24 @@ let autodiff_level_rev_transfer let gen = Map.find_exn fwd_ad_levels l in let kill = match mir_node with - | Decl {decl_id; _} -> Set.Poly.singleton decl_id + | Decl { decl_id; _ } -> Set.Poly.singleton decl_id | _ -> Set.Poly.empty in transfer_gen_kill_alt p gen kill + ;; (* gens and then kills *) - end - : TRANSFER_FUNCTION - with type labels = int and type properties = string Set.Poly.t ) + end : TRANSFER_FUNCTION + with type labels = int + and type properties = string Set.Poly.t) +;; (** The transfer function for the second forward analysis part of determining optimal ad-levels for variables *) let autodiff_level_fwd2_transfer (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (rev_ad_levels : (int, string Set.Poly.t) Map.Poly.t) = - ( module struct + (rev_ad_levels : (int, string Set.Poly.t) Map.Poly.t) + = + (module struct type labels = int type properties = string Set.Poly.t @@ -823,13 +885,15 @@ let autodiff_level_fwd2_transfer let gen = Map.find_exn rev_ad_levels l in let kill = match mir_node with - | Decl {decl_id; _} -> Set.Poly.singleton decl_id + | Decl { decl_id; _ } -> Set.Poly.singleton decl_id | _ -> Set.Poly.empty in transfer_gen_kill p gen kill - end - : TRANSFER_FUNCTION - with type labels = int and type properties = string Set.Poly.t ) + ;; + end : TRANSFER_FUNCTION + with type labels = int + and type properties = string Set.Poly.t) +;; (** The central definition of a monotone dataflow analysis framework. Given a compatible flowgraph, lattice and transfer function, we can @@ -840,11 +904,13 @@ let autodiff_level_fwd2_transfer solution that we would really be interested in, but which is often incomputable. In case of a distributive lattice of properties, the MFP and MOP solutions coincide. *) -let monotone_framework (type l p) (module F : FLOWGRAPH with type labels = l) +let monotone_framework + (type l p) + (module F : FLOWGRAPH with type labels = l) (module L : LATTICE with type properties = p) (module T : TRANSFER_FUNCTION with type labels = l and type properties = p) - = - ( module struct + = + (module struct type labels = l type properties = p @@ -853,77 +919,83 @@ let monotone_framework (type l p) (module F : FLOWGRAPH with type labels = l) let workstack = Stack.create () in (* TODO: does the order matter a lot for efficiency here? *) Map.iteri F.successors ~f:(fun ~key ~data -> - Set.iter data ~f:(fun succ -> Stack.push workstack (key, succ)) ) ; + Set.iter data ~f:(fun succ -> Stack.push workstack (key, succ))); let analysis_in = Hashtbl.create (module F) in Map.iter_keys ~f:(fun l -> - Hashtbl.add_exn analysis_in ~key:l - ~data:(if Set.mem F.initials l then L.initial else L.bottom) ) - F.successors ; + Hashtbl.add_exn + analysis_in + ~key:l + ~data:(if Set.mem F.initials l then L.initial else L.bottom)) + F.successors; (* STEP 2: iterate *) while Stack.length workstack <> 0 do let l, l' = Stack.pop_exn workstack in let old_analysis_in_l' = Hashtbl.find_exn analysis_in l' in - let new_analysis_in_l' = - T.transfer_function l (Hashtbl.find_exn analysis_in l) - in - if not (L.leq new_analysis_in_l' old_analysis_in_l') then + let new_analysis_in_l' = T.transfer_function l (Hashtbl.find_exn analysis_in l) in + if not (L.leq new_analysis_in_l' old_analysis_in_l') + then ( let () = - Hashtbl.set analysis_in ~key:l' + Hashtbl.set + analysis_in + ~key:l' ~data:(L.lub old_analysis_in_l' new_analysis_in_l') in Set.iter (Map.find_exn F.successors l') ~f:(fun l'' -> - Stack.push workstack (l', l'') ) - done ; + Stack.push workstack (l', l''))) + done; (* STEP 3: present final results *) let analysis_in_out = - Map.fold ~init:Map.Poly.empty + Map.fold + ~init:Map.Poly.empty ~f:(fun ~key ~data:_ accum -> let analysis_in_data = Hashtbl.find_exn analysis_in key in - Map.add_exn accum ~key + Map.add_exn + accum + ~key ~data: - { entry= analysis_in_data - ; exit= T.transfer_function key analysis_in_data } ) + { entry = analysis_in_data + ; exit = T.transfer_function key analysis_in_data + }) F.successors in analysis_in_out - end - : MONOTONE_FRAMEWORK - with type labels = l and type properties = p ) + ;; + end : MONOTONE_FRAMEWORK + with type labels = l + and type properties = p) +;; -let rec declared_variables_stmt - (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = +let rec declared_variables_stmt (s : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = match s with - | Decl {decl_id= x; _} -> Set.Poly.singleton x + | Decl { decl_id = x; _ } -> Set.Poly.singleton x | Assignment (_, _) - |TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip -> - Set.Poly.empty + | TargetPE _ + | NRFunApp (_, _, _) + | Break | Continue | Return _ | Skip -> Set.Poly.empty | IfElse (_, b1, Some b2) -> - Set.Poly.union - (declared_variables_stmt b1.pattern) - (declared_variables_stmt b2.pattern) + Set.Poly.union + (declared_variables_stmt b1.pattern) + (declared_variables_stmt b2.pattern) | While (_, b) | IfElse (_, b, None) -> declared_variables_stmt b.pattern - | For {loopvar= s; body= b; _} -> - Set.Poly.add (declared_variables_stmt b.pattern) s + | For { loopvar = s; body = b; _ } -> Set.Poly.add (declared_variables_stmt b.pattern) s | Block l | SList l -> - Set.Poly.union_list - (List.map ~f:(fun x -> declared_variables_stmt x.pattern) l) + Set.Poly.union_list (List.map ~f:(fun x -> declared_variables_stmt x.pattern) l) +;; -let propagation_mfp (prog : Program.Typed.t) - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) +let propagation_mfp + (prog : Program.Typed.t) + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) (propagation_transfer : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t - -> (module - TRANSFER_FUNCTION + (int, Stmt.Located.Non_recursive.t) Map.Poly.t + -> (module TRANSFER_FUNCTION with type labels = int - and type properties = (string, Expr.Typed.t) Map.Poly.t option)) = + and type properties = (string, Expr.Typed.t) Map.Poly.t option)) + = let mir = Map.find_exn flowgraph_to_mir 1 in let domain = - ( module struct + (module struct type vals = string let total = @@ -931,44 +1003,49 @@ let propagation_mfp (prog : Program.Typed.t) [ Set.Poly.of_list (List.map ~f:fst prog.input_vars) ; Set.Poly.of_list (List.map ~f:fst prog.output_vars) ; declared_variables_stmt - (stmt_loc_of_stmt_loc_num flowgraph_to_mir mir).pattern ] - end - : TOTALTYPE - with type vals = string ) + (stmt_loc_of_stmt_loc_num flowgraph_to_mir mir).pattern + ] + ;; + end : TOTALTYPE + with type vals = string) in let codomain = - (module struct type vals = Expr.Typed.t - end - : TYPE - with type vals = Expr.Typed.t ) - in - let (module Lattice) = - dual_partial_function_lattice_with_bot domain codomain + (module struct + type vals = Expr.Typed.t + end : TYPE + with type vals = Expr.Typed.t) in + let (module Lattice) = dual_partial_function_lattice_with_bot domain codomain in let (module Transfer) = propagation_transfer flowgraph_to_mir in let (module Mf) = monotone_framework (module Flowgraph) (module Lattice) (module Transfer) in Mf.mfp () +;; -let reaching_definitions_mfp (mir : Program.Typed.t) - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = +let reaching_definitions_mfp + (mir : Program.Typed.t) + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = let variables = - ( module struct + (module struct type vals = string let initial = Set.Poly.union_list [ Set.Poly.of_list (List.map ~f:fst mir.input_vars) - ; Set.Poly.of_list (List.map ~f:fst mir.output_vars) ] - end - : INITIALTYPE - with type vals = string ) + ; Set.Poly.of_list (List.map ~f:fst mir.output_vars) + ] + ;; + end : INITIALTYPE + with type vals = string) in let labels = - (module struct type vals = int end : TYPE with type vals = int) + (module struct + type vals = int + end : TYPE + with type vals = int) in let (module Lattice) = reaching_definitions_lattice variables labels in let (module Transfer) = reaching_definitions_transfer flowgraph_to_mir in @@ -976,76 +1053,80 @@ let reaching_definitions_mfp (mir : Program.Typed.t) monotone_framework (module Flowgraph) (module Lattice) (module Transfer) in Mf.mfp () +;; -let initialized_vars_mfp (total : string Set.Poly.t) - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = +let initialized_vars_mfp + (total : string Set.Poly.t) + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = let (module Lattice) = dual_powerset_lattice_empty_initial - (module struct type vals = string + (module struct + type vals = string - let total = total end) + let total = total + end) in let (module Transfer) = initialized_vars_transfer flowgraph_to_mir in let (module Mf) = monotone_framework (module Flowgraph) (module Lattice) (module Transfer) in Mf.mfp () +;; let globals (prog : Program.Typed.t) = Set.Poly.union_list [ Set.Poly.of_list (List.map ~f:fst prog.output_vars) (* It is not strictly necessary to exclude data variables from DCE. - However, - 1. We don't currently check for usage of data variables in - corners of the MIR, such as in the sizes of parameters - 2. There is code added in codegen that is never represented in - the MIR that may use data variables as if they're initialized - *) + However, + 1. We don't currently check for usage of data variables in + corners of the MIR, such as in the sizes of parameters + 2. There is code added in codegen that is never represented in + the MIR that may use data variables as if they're initialized + *) ; Set.Poly.of_list (List.map ~f:fst prog.input_vars) ; Set.Poly.union_list (List.map ~f:var_declarations prog.prepare_data) - ; Set.Poly.singleton "target" ] + ; Set.Poly.singleton "target" + ] +;; (** Monotone framework instance for live_variables analysis. Expects reverse flowgraph. *) -let live_variables_mfp (prog : Program.Typed.t) - (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = +let live_variables_mfp + (prog : Program.Typed.t) + (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = let never_kill = globals prog in let variables = - ( module struct + (module struct type vals = string (* NOTE: global generated quantities, (transformed) parameters and target are always observable - so should be live. *) + so should be live. *) let initial = never_kill - end - : INITIALTYPE - with type vals = string ) + end : INITIALTYPE + with type vals = string) in let (module Lattice) = powerset_lattice variables in - let (module Transfer) = - live_variables_transfer never_kill flowgraph_to_mir - in + let (module Transfer) = live_variables_transfer never_kill flowgraph_to_mir in let (module Mf) = monotone_framework (module Rev_Flowgraph) (module Lattice) (module Transfer) in Mf.mfp () +;; (** Instantiate all four instances of the monotone framework for lazy code motion, reusing code between them *) let lazy_expressions_mfp - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) = + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) + = let all_expressions = used_subexpressions_stmt - (stmt_loc_of_stmt_loc_num flowgraph_to_mir - (Map.Poly.find_exn flowgraph_to_mir 1)) + (stmt_loc_of_stmt_loc_num flowgraph_to_mir (Map.Poly.find_exn flowgraph_to_mir 1)) .pattern in (* TODO: this could probably be done in a nicer way *) @@ -1054,14 +1135,9 @@ let lazy_expressions_mfp dual_powerset_lattice_expressions Expr.Typed.Set.empty all_expressions in let (module Lattice2) = powerset_lattice_expressions Expr.Typed.Set.empty in - let (module Transfer1) = - anticipated_expressions_transfer flowgraph_to_mir used_expr - in + let (module Transfer1) = anticipated_expressions_transfer flowgraph_to_mir used_expr in let (module Mf1) = - monotone_framework - (module Rev_Flowgraph) - (module Lattice1) - (module Transfer1) + monotone_framework (module Rev_Flowgraph) (module Lattice1) (module Transfer1) in let anticipated_expressions_mfp = Mf1.mfp () in let (module Transfer2) = @@ -1071,40 +1147,30 @@ let lazy_expressions_mfp monotone_framework (module Flowgraph) (module Lattice1) (module Transfer2) in let available_expressions_mfp = Mf2.mfp () in - let earliest_expr = - earliest anticipated_expressions_mfp available_expressions_mfp - in - let (module Transfer3) = - postponable_expressions_transfer earliest_expr used_expr - in + let earliest_expr = earliest anticipated_expressions_mfp available_expressions_mfp in + let (module Transfer3) = postponable_expressions_transfer earliest_expr used_expr in let (module Mf3) = monotone_framework (module Flowgraph) (module Lattice1) (module Transfer3) in let postponable_expressions_mfp = Mf3.mfp () in let latest_expr = - latest Flowgraph.successors earliest_expr postponable_expressions_mfp - used_expr - in - let (module Transfer4) = - used_not_latest_expressions_transfer used_expr latest_expr + latest Flowgraph.successors earliest_expr postponable_expressions_mfp used_expr in + let (module Transfer4) = used_not_latest_expressions_transfer used_expr latest_expr in let (module Mf4) = - monotone_framework - (module Rev_Flowgraph) - (module Lattice2) - (module Transfer4) + monotone_framework (module Rev_Flowgraph) (module Lattice2) (module Transfer4) in let used_not_latest_expressions_mfp = Mf4.mfp () in - (latest_expr, used_not_latest_expressions_mfp) + latest_expr, used_not_latest_expressions_mfp +;; (** Perform the analysis for ad-levels, using both the fwd and reverse pass *) let autodiff_level_mfp - (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) - (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH - with type labels = int) + (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) + (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int) (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) - (autodiff_variables : string Set.Poly.t) = + (autodiff_variables : string Set.Poly.t) + = let (module Lattice1) = autodiff_level_lattice autodiff_variables in let (module Lattice2) = autodiff_level_lattice Set.Poly.empty in let (module Transfer1) = autodiff_level_fwd1_transfer flowgraph_to_mir in @@ -1113,18 +1179,17 @@ let autodiff_level_mfp in let fwd1_ad_levels_mfp = Mf1.mfp () in let (module Transfer2) = - autodiff_level_rev_transfer flowgraph_to_mir + autodiff_level_rev_transfer + flowgraph_to_mir (Map.map ~f:(fun x -> x.exit) fwd1_ad_levels_mfp) in let (module Mf2) = - monotone_framework - (module Rev_Flowgraph) - (module Lattice2) - (module Transfer2) + monotone_framework (module Rev_Flowgraph) (module Lattice2) (module Transfer2) in let rev_ad_levels_mfp = Mf2.mfp () in let (module Transfer3) = - autodiff_level_fwd2_transfer flowgraph_to_mir + autodiff_level_fwd2_transfer + flowgraph_to_mir (Map.map ~f:(fun x -> x.entry) rev_ad_levels_mfp) in let (module Mf3) = @@ -1132,3 +1197,4 @@ let autodiff_level_mfp in let fwd2_ad_levels_mfp = Mf3.mfp () in fwd2_ad_levels_mfp +;; diff --git a/src/analysis_and_optimization/Monotone_framework_sigs.mli b/src/analysis_and_optimization/Monotone_framework_sigs.mli index 3aefb4d740..51bc1bf3ab 100644 --- a/src/analysis_and_optimization/Monotone_framework_sigs.mli +++ b/src/analysis_and_optimization/Monotone_framework_sigs.mli @@ -46,10 +46,10 @@ module type LATTICE_NO_BOT = sig val leq : properties -> properties -> bool - val initial : properties (** An initial value, which might not be the top element. The idea is that this is the property that you start with (you assume to be true at the start of your analysis). *) + val initial : properties val lub : properties -> properties -> properties end @@ -73,7 +73,10 @@ module type TRANSFER_FUNCTION = sig val transfer_function : labels -> properties -> properties end -type 'a entry_exit = {entry: 'a; exit: 'a} +type 'a entry_exit = + { entry : 'a + ; exit : 'a + } (** The API for a monotone framework. mfp computes the minimal fixed point of the equations/inequalities defined between property lattice diff --git a/src/analysis_and_optimization/Optimize.ml b/src/analysis_and_optimization/Optimize.ml index 1ef0ecfd24..58d371726b 100644 --- a/src/analysis_and_optimization/Optimize.ml +++ b/src/analysis_and_optimization/Optimize.ml @@ -10,54 +10,62 @@ let preserve_stability = false Apply the transformation to each function body and to the rest of the program as one block. *) -let transform_program (mir : Program.Typed.t) - (transform : Stmt.Located.t -> Stmt.Located.t) : Program.Typed.t = +let transform_program + (mir : Program.Typed.t) + (transform : Stmt.Located.t -> Stmt.Located.t) + : Program.Typed.t + = let packed_prog_body = transform - { pattern= + { pattern = SList (List.map - ~f:(fun x -> - Stmt.Fixed.{pattern= SList x; meta= Location_span.empty} ) - [ mir.prepare_data; mir.transform_inits; mir.log_prob - ; mir.generate_quantities ]) - ; meta= Location_span.empty } + ~f:(fun x -> Stmt.Fixed.{ pattern = SList x; meta = Location_span.empty }) + [ mir.prepare_data + ; mir.transform_inits + ; mir.log_prob + ; mir.generate_quantities + ]) + ; meta = Location_span.empty + } in let transformed_prog_body = transform packed_prog_body in let transformed_functions = - List.map mir.functions_block ~f:(fun fs -> - {fs with fdbody= transform fs.fdbody} ) + List.map mir.functions_block ~f:(fun fs -> { fs with fdbody = transform fs.fdbody }) in match transformed_prog_body with - | { pattern= + | { pattern = SList - [ {pattern= SList prepare_data'; _} - ; {pattern= SList transform_inits'; _} - ; {pattern= SList log_prob'; _} - ; {pattern= SList generate_quantities'; _} ]; _ } -> - { mir with - functions_block= transformed_functions - ; prepare_data= prepare_data' - ; transform_inits= transform_inits' - ; log_prob= log_prob' - ; generate_quantities= generate_quantities' } - | _ -> - raise - (Failure "Something went wrong with program transformation packing!") + [ { pattern = SList prepare_data'; _ } + ; { pattern = SList transform_inits'; _ } + ; { pattern = SList log_prob'; _ } + ; { pattern = SList generate_quantities'; _ } + ] + ; _ + } -> + { mir with + functions_block = transformed_functions + ; prepare_data = prepare_data' + ; transform_inits = transform_inits' + ; log_prob = log_prob' + ; generate_quantities = generate_quantities' + } + | _ -> raise (Failure "Something went wrong with program transformation packing!") +;; (** Apply the transformation to each function body and to each program block separately. *) -let transform_program_blockwise (mir : Program.Typed.t) +let transform_program_blockwise + (mir : Program.Typed.t) (transform : Stmt.Located.t Program.fun_def option -> Stmt.Located.t -> Stmt.Located.t) - : Program.Typed.t = + : Program.Typed.t + = let transform' fd s = - match transform fd {pattern= SList s; meta= Location_span.empty} with - | {pattern= SList l; _} -> l - | _ -> - raise - (Failure "Something went wrong with program transformation packing!") + match transform fd { pattern = SList s; meta = Location_span.empty } with + | { pattern = SList l; _ } -> l + | _ -> raise (Failure "Something went wrong with program transformation packing!") in (* Right now, we have an implicit constraint where if fdbody = Skip, the fun_def is a function declaration. When that's the case we don't want @@ -67,57 +75,65 @@ let transform_program_blockwise (mir : Program.Typed.t) in let transformed_functions = List.map non_decl_functions ~f:(fun fs -> - {fs with fdbody= transform (Some fs) fs.fdbody} ) + { fs with fdbody = transform (Some fs) fs.fdbody }) in { mir with - functions_block= transformed_functions - ; prepare_data= transform' None mir.prepare_data - ; transform_inits= transform' None mir.transform_inits - ; log_prob= transform' None mir.log_prob - ; generate_quantities= transform' None mir.generate_quantities } + functions_block = transformed_functions + ; prepare_data = transform' None mir.prepare_data + ; transform_inits = transform' None mir.transform_inits + ; log_prob = transform' None mir.log_prob + ; generate_quantities = transform' None mir.generate_quantities + } +;; let map_no_loc l = - List.map ~f:(fun s -> Stmt.Fixed.{pattern= s; meta= Location_span.empty}) l + List.map ~f:(fun s -> Stmt.Fixed.{ pattern = s; meta = Location_span.empty }) l +;; let slist_no_loc l = Stmt.Fixed.Pattern.SList (map_no_loc l) let block_no_loc l = Stmt.Fixed.Pattern.Block (map_no_loc l) let slist_concat_no_loc l stmt = - match l with [] -> stmt | l -> slist_no_loc (l @ [stmt]) + match l with + | [] -> stmt + | l -> slist_no_loc (l @ [ stmt ]) +;; let replace_fresh_local_vars s' = let f m = function - | Stmt.Fixed.Pattern.Decl {decl_adtype; decl_type; decl_id} -> - let new_name = - match Map.Poly.find m decl_id with - | Some existing -> existing - | None -> Gensym.generate ~prefix:"inline_" () - in - ( Stmt.Fixed.Pattern.Decl {decl_adtype; decl_id= new_name; decl_type} - , Map.Poly.set m ~key:decl_id ~data:new_name ) - | Stmt.Fixed.Pattern.For {loopvar; lower; upper; body} -> - let new_name = - match Map.Poly.find m loopvar with - | Some existing -> existing - | None -> Gensym.generate ~prefix:"inline_" () - in - ( Stmt.Fixed.Pattern.For {loopvar= new_name; lower; upper; body} - , Map.Poly.set m ~key:loopvar ~data:new_name ) + | Stmt.Fixed.Pattern.Decl { decl_adtype; decl_type; decl_id } -> + let new_name = + match Map.Poly.find m decl_id with + | Some existing -> existing + | None -> Gensym.generate ~prefix:"inline_" () + in + ( Stmt.Fixed.Pattern.Decl { decl_adtype; decl_id = new_name; decl_type } + , Map.Poly.set m ~key:decl_id ~data:new_name ) + | Stmt.Fixed.Pattern.For { loopvar; lower; upper; body } -> + let new_name = + match Map.Poly.find m loopvar with + | Some existing -> existing + | None -> Gensym.generate ~prefix:"inline_" () + in + ( Stmt.Fixed.Pattern.For { loopvar = new_name; lower; upper; body } + , Map.Poly.set m ~key:loopvar ~data:new_name ) | Assignment ((var_name, ut, l), e) -> - let var_name = - match Map.Poly.find m var_name with - | None -> var_name - | Some var_name -> var_name - in - (Stmt.Fixed.Pattern.Assignment ((var_name, ut, l), e), m) - | x -> (x, m) + let var_name = + match Map.Poly.find m var_name with + | None -> var_name + | Some var_name -> var_name + in + Stmt.Fixed.Pattern.Assignment ((var_name, ut, l), e), m + | x -> x, m in let s, m = map_rec_state_stmt_loc f Map.Poly.empty s' in name_subst_stmt m s +;; let subst_args_stmt args es = let m = Map.Poly.of_alist_exn (List.zip_exn args es) in subst_stmt m +;; (* TODO: only handle early returns if that's necessary *) (* The strategy here is to wrap the function body in a dummy loop, then replace @@ -130,301 +146,347 @@ let subst_args_stmt args es = let handle_early_returns opt_var b = let returned = Gensym.generate ~prefix:"inline_" () in let f = function - | Stmt.Fixed.Pattern.Return opt_ret -> ( - match (opt_var, opt_ret) with + | Stmt.Fixed.Pattern.Return opt_ret -> + (match opt_var, opt_ret with | None, None -> Stmt.Fixed.Pattern.Break | Some name, Some e -> - SList - [ Stmt.Fixed. - { pattern= - Assignment - ( (returned, UInt, []) - , Expr.Fixed. - { pattern= Lit (Int, "1") - ; meta= - Expr.Typed.Meta. - { type_= UInt - ; adlevel= DataOnly - ; loc= Location_span.empty } } ) - ; meta= Location_span.empty } - ; Stmt.Fixed. - { pattern= Assignment ((name, Expr.Typed.type_of e, []), e) - ; meta= Location_span.empty } - ; {pattern= Break; meta= Location_span.empty} ] + SList + [ Stmt.Fixed. + { pattern = + Assignment + ( (returned, UInt, []) + , Expr.Fixed. + { pattern = Lit (Int, "1") + ; meta = + Expr.Typed.Meta. + { type_ = UInt + ; adlevel = DataOnly + ; loc = Location_span.empty + } + } ) + ; meta = Location_span.empty + } + ; Stmt.Fixed. + { pattern = Assignment ((name, Expr.Typed.type_of e, []), e) + ; meta = Location_span.empty + } + ; { pattern = Break; meta = Location_span.empty } + ] | Some _, None -> - raise_s - [%message - ( "Function should return a value but found an empty return \ - statement." - : string )] + raise_s + [%message + ("Function should return a value but found an empty return statement." + : string)] | None, Some _ -> - raise_s - [%message - ( "Expected a void function but found a non-empty return \ - statement." - : string )] ) + raise_s + [%message + ("Expected a void function but found a non-empty return statement." : string)]) | Stmt.Fixed.Pattern.For _ as loop -> - Stmt.Fixed.Pattern.SList - [ Stmt.Fixed.{pattern= loop; meta= Location_span.empty} - ; Stmt.Fixed. - { pattern= - IfElse - ( Expr.Fixed. - { pattern= Var returned - ; meta= - Expr.Typed.Meta. - { type_= UInt - ; adlevel= DataOnly - ; loc= Location_span.empty } } - , {pattern= Break; meta= Location_span.empty} - , None ) - ; meta= Location_span.empty } ] + Stmt.Fixed.Pattern.SList + [ Stmt.Fixed.{ pattern = loop; meta = Location_span.empty } + ; Stmt.Fixed. + { pattern = + IfElse + ( Expr.Fixed. + { pattern = Var returned + ; meta = + Expr.Typed.Meta. + { type_ = UInt + ; adlevel = DataOnly + ; loc = Location_span.empty + } + } + , { pattern = Break; meta = Location_span.empty } + , None ) + ; meta = Location_span.empty + } + ] | x -> x in Stmt.Fixed.Pattern.SList [ Stmt.Fixed. - { pattern= - Decl - {decl_adtype= DataOnly; decl_id= returned; decl_type= Sized SInt} - ; meta= Location_span.empty } + { pattern = + Decl { decl_adtype = DataOnly; decl_id = returned; decl_type = Sized SInt } + ; meta = Location_span.empty + } ; Stmt.Fixed. - { pattern= + { pattern = Assignment ( (returned, UInt, []) , Expr.Fixed. - { pattern= Lit (Int, "0") - ; meta= + { pattern = Lit (Int, "0") + ; meta = Expr.Typed.Meta. - { type_= UInt - ; adlevel= DataOnly - ; loc= Location_span.empty } } ) - ; meta= Location_span.empty } + { type_ = UInt; adlevel = DataOnly; loc = Location_span.empty } + } ) + ; meta = Location_span.empty + } ; Stmt.Fixed. - { pattern= + { pattern = Stmt.Fixed.Pattern.For - { loopvar= Gensym.generate ~prefix:"inline_" () - ; lower= + { loopvar = Gensym.generate ~prefix:"inline_" () + ; lower = Expr.Fixed. - { pattern= Lit (Int, "1") - ; meta= + { pattern = Lit (Int, "1") + ; meta = Expr.Typed.Meta. - { type_= UInt - ; adlevel= DataOnly - ; loc= Location_span.empty } } - ; upper= - { pattern= Lit (Int, "1") - ; meta= - {type_= UInt; adlevel= DataOnly; loc= Location_span.empty} + { type_ = UInt; adlevel = DataOnly; loc = Location_span.empty } + } + ; upper = + { pattern = Lit (Int, "1") + ; meta = { type_ = UInt; adlevel = DataOnly; loc = Location_span.empty } } - ; body= map_rec_stmt_loc f b } - ; meta= Location_span.empty } ] + ; body = map_rec_stmt_loc f b + } + ; meta = Location_span.empty + } + ] +;; (* Triple is (declaration list, statement list, return expression) *) -let rec inline_function_expression adt fim (Expr.Fixed.({pattern; _}) as e) = +let rec inline_function_expression adt fim (Expr.Fixed.{ pattern; _ } as e) = match pattern with - | Var _ -> ([], [], e) - | Lit (_, _) -> ([], [], e) - | FunApp (t, s, es) -> ( - let dse_list = List.map ~f:(inline_function_expression adt fim) es in - (* function arguments are evaluated from right to left in C++, so we need to reverse *) - let d_list = - List.concat (List.rev (List.map ~f:(function x, _, _ -> x) dse_list)) - in - let s_list = - List.concat (List.rev (List.map ~f:(function _, x, _ -> x) dse_list)) + | Var _ -> [], [], e + | Lit (_, _) -> [], [], e + | FunApp (t, s, es) -> + let dse_list = List.map ~f:(inline_function_expression adt fim) es in + (* function arguments are evaluated from right to left in C++, so we need to reverse *) + let d_list = + List.concat + (List.rev + (List.map + ~f:(function + | x, _, _ -> x) + dse_list)) + in + let s_list = + List.concat + (List.rev + (List.map + ~f:(function + | _, x, _ -> x) + dse_list)) + in + let es = + List.map + ~f:(function + | _, _, x -> x) + dse_list + in + (match Map.find fim s with + | None -> d_list, s_list, { e with pattern = FunApp (t, s, es) } + | Some (rt, args, b) -> + let x = Gensym.generate ~prefix:"inline_" () in + let handle = handle_early_returns (Some x) in + let d_list2, s_list2, (e : Expr.Typed.t) = + ( [ Stmt.Fixed.Pattern.Decl + { decl_adtype = adt; decl_id = x; decl_type = Option.value_exn rt } + ] + (* We should minimize the code that's having its variables + replaced to avoid conflict with the (two) new dummy + variables introduced by inlining *) + , [ handle (replace_fresh_local_vars (subst_args_stmt args es b)) ] + , { pattern = Var x + ; meta = + Expr.Typed.Meta. + { type_ = Type.to_unsized (Option.value_exn rt) + ; adlevel = adt + ; loc = Location_span.empty + } + } ) in - let es = List.map ~f:(function _, _, x -> x) dse_list in - match Map.find fim s with - | None -> (d_list, s_list, {e with pattern= FunApp (t, s, es)}) - | Some (rt, args, b) -> - let x = Gensym.generate ~prefix:"inline_" () in - let handle = handle_early_returns (Some x) in - let d_list2, s_list2, (e : Expr.Typed.t) = - ( [ Stmt.Fixed.Pattern.Decl - {decl_adtype= adt; decl_id= x; decl_type= Option.value_exn rt} - ] - (* We should minimize the code that's having its variables - replaced to avoid conflict with the (two) new dummy - variables introduced by inlining *) - , [handle (replace_fresh_local_vars (subst_args_stmt args es b))] - , { pattern= Var x - ; meta= - Expr.Typed.Meta. - { type_= Type.to_unsized (Option.value_exn rt) - ; adlevel= adt - ; loc= Location_span.empty } } ) - in - let d_list = d_list @ d_list2 in - let s_list = s_list @ s_list2 in - (d_list, s_list, e) ) + let d_list = d_list @ d_list2 in + let s_list = s_list @ s_list2 in + d_list, s_list, e) | TernaryIf (e1, e2, e3) -> - let dl1, sl1, e1 = inline_function_expression adt fim e1 in - let dl2, sl2, e2 = inline_function_expression adt fim e2 in - let dl3, sl3, e3 = inline_function_expression adt fim e3 in - ( dl1 @ dl2 @ dl3 - , sl1 - @ [ Stmt.Fixed.( - Pattern.IfElse - ( e1 - , {pattern= block_no_loc sl2; meta= Location_span.empty} - , Some {pattern= block_no_loc sl3; meta= Location_span.empty} - )) ] - , {e with pattern= TernaryIf (e1, e2, e3)} ) - | Indexed (e', i_list) -> - let dl, sl, e' = inline_function_expression adt fim e' in - let dsi_list = List.map ~f:(inline_function_index adt fim) i_list in - let d_list = - List.concat (List.rev (List.map ~f:(function x, _, _ -> x) dsi_list)) - in - let s_list = - List.concat (List.rev (List.map ~f:(function _, x, _ -> x) dsi_list)) - in - let i_list = List.map ~f:(function _, _, x -> x) dsi_list in - (d_list @ dl, s_list @ sl, {e with pattern= Indexed (e', i_list)}) - | EAnd (e1, e2) -> - let dl1, sl1, e1 = inline_function_expression adt fim e1 in - let dl2, sl2, e2 = inline_function_expression adt fim e2 in - let sl2 = - [ Stmt.Fixed.( + let dl1, sl1, e1 = inline_function_expression adt fim e1 in + let dl2, sl2, e2 = inline_function_expression adt fim e2 in + let dl3, sl3, e3 = inline_function_expression adt fim e3 in + ( dl1 @ dl2 @ dl3 + , sl1 + @ [ Stmt.Fixed.( Pattern.IfElse ( e1 - , {pattern= Block (map_no_loc sl2); meta= Location_span.empty} - , None )) ] - in - (dl1 @ dl2, sl1 @ sl2, {e with pattern= EAnd (e1, e2)}) + , { pattern = block_no_loc sl2; meta = Location_span.empty } + , Some { pattern = block_no_loc sl3; meta = Location_span.empty } )) + ] + , { e with pattern = TernaryIf (e1, e2, e3) } ) + | Indexed (e', i_list) -> + let dl, sl, e' = inline_function_expression adt fim e' in + let dsi_list = List.map ~f:(inline_function_index adt fim) i_list in + let d_list = + List.concat + (List.rev + (List.map + ~f:(function + | x, _, _ -> x) + dsi_list)) + in + let s_list = + List.concat + (List.rev + (List.map + ~f:(function + | _, x, _ -> x) + dsi_list)) + in + let i_list = + List.map + ~f:(function + | _, _, x -> x) + dsi_list + in + d_list @ dl, s_list @ sl, { e with pattern = Indexed (e', i_list) } + | EAnd (e1, e2) -> + let dl1, sl1, e1 = inline_function_expression adt fim e1 in + let dl2, sl2, e2 = inline_function_expression adt fim e2 in + let sl2 = + [ Stmt.Fixed.( + Pattern.IfElse + (e1, { pattern = Block (map_no_loc sl2); meta = Location_span.empty }, None)) + ] + in + dl1 @ dl2, sl1 @ sl2, { e with pattern = EAnd (e1, e2) } | EOr (e1, e2) -> - let dl1, sl1, e1 = inline_function_expression adt fim e1 in - let dl2, sl2, e2 = inline_function_expression adt fim e2 in - let sl2 = - [ Stmt.Fixed.( - Pattern.IfElse - ( e1 - , {pattern= Skip; meta= Location_span.empty} - , Some - {pattern= Block (map_no_loc sl2); meta= Location_span.empty} - )) ] - in - (dl1 @ dl2, sl1 @ sl2, {e with pattern= EOr (e1, e2)}) + let dl1, sl1, e1 = inline_function_expression adt fim e1 in + let dl2, sl2, e2 = inline_function_expression adt fim e2 in + let sl2 = + [ Stmt.Fixed.( + Pattern.IfElse + ( e1 + , { pattern = Skip; meta = Location_span.empty } + , Some { pattern = Block (map_no_loc sl2); meta = Location_span.empty } )) + ] + in + dl1 @ dl2, sl1 @ sl2, { e with pattern = EOr (e1, e2) } and inline_function_index adt fim i = match i with - | All -> ([], [], All) + | All -> [], [], All | Single e -> - let dl, sl, e = inline_function_expression adt fim e in - (dl, sl, Single e) + let dl, sl, e = inline_function_expression adt fim e in + dl, sl, Single e | Upfrom e -> - let dl, sl, e = inline_function_expression adt fim e in - (dl, sl, Upfrom e) + let dl, sl, e = inline_function_expression adt fim e in + dl, sl, Upfrom e | Between (e1, e2) -> - let dl1, sl1, e1 = inline_function_expression adt fim e1 in - let dl2, sl2, e2 = inline_function_expression adt fim e2 in - (dl1 @ dl2, sl1 @ sl2, Between (e1, e2)) + let dl1, sl1, e1 = inline_function_expression adt fim e1 in + let dl2, sl2, e2 = inline_function_expression adt fim e2 in + dl1 @ dl2, sl1 @ sl2, Between (e1, e2) | MultiIndex e -> - let dl, sl, e = inline_function_expression adt fim e in - (dl, sl, MultiIndex e) + let dl, sl, e = inline_function_expression adt fim e in + dl, sl, MultiIndex e +;; -let rec inline_function_statement adt fim Stmt.Fixed.({pattern; meta}) = +let rec inline_function_statement adt fim Stmt.Fixed.{ pattern; meta } = Stmt.Fixed. - { pattern= - ( match pattern with + { pattern = + (match pattern with | Assignment ((x, ut, l), e2) -> - let e1 = - {e2 with pattern= Indexed ({e2 with pattern= Var x}, l)} - in - (* This inner e2 is wrong. We are giving the wrong type to Var x. But it doens't really matter as we discard it later. *) - let dl1, sl1, e1 = inline_function_expression adt fim e1 in - let dl2, sl2, e2 = inline_function_expression adt fim e2 in - let x, l = - match e1.pattern with - | Var x -> (x, []) - | Indexed ({pattern= Var x; _}, l) -> (x, l) - | _ as w -> - raise_s [%sexp (w : Expr.Typed.t Expr.Fixed.Pattern.t)] - in - slist_concat_no_loc - (dl2 @ dl1 @ sl2 @ sl1) - (Assignment ((x, ut, l), e2)) + let e1 = { e2 with pattern = Indexed ({ e2 with pattern = Var x }, l) } in + (* This inner e2 is wrong. We are giving the wrong type to Var x. But it doens't really matter as we discard it later. *) + let dl1, sl1, e1 = inline_function_expression adt fim e1 in + let dl2, sl2, e2 = inline_function_expression adt fim e2 in + let x, l = + match e1.pattern with + | Var x -> x, [] + | Indexed ({ pattern = Var x; _ }, l) -> x, l + | _ as w -> raise_s [%sexp (w : Expr.Typed.t Expr.Fixed.Pattern.t)] + in + slist_concat_no_loc (dl2 @ dl1 @ sl2 @ sl1) (Assignment ((x, ut, l), e2)) | TargetPE e -> - let d, s, e = inline_function_expression adt fim e in - slist_concat_no_loc (d @ s) (TargetPE e) + let d, s, e = inline_function_expression adt fim e in + slist_concat_no_loc (d @ s) (TargetPE e) | NRFunApp (t, s, es) -> - let dse_list = - List.map ~f:(inline_function_expression adt fim) es - in - (* function arguments are evaluated from right to left in C++, so we need to reverse *) - let d_list = - List.concat - (List.rev (List.map ~f:(function x, _, _ -> x) dse_list)) - in - let s_list = - List.concat - (List.rev (List.map ~f:(function _, x, _ -> x) dse_list)) - in - let es = List.map ~f:(function _, _, x -> x) dse_list in - slist_concat_no_loc (d_list @ s_list) - ( match Map.find fim s with - | None -> NRFunApp (t, s, es) - | Some (_, args, b) -> - let b = replace_fresh_local_vars b in - let b = handle_early_returns None b in - (subst_args_stmt args es - {pattern= b; meta= Location_span.empty}) - .pattern ) - | Return e -> ( - match e with + let dse_list = List.map ~f:(inline_function_expression adt fim) es in + (* function arguments are evaluated from right to left in C++, so we need to reverse *) + let d_list = + List.concat + (List.rev + (List.map + ~f:(function + | x, _, _ -> x) + dse_list)) + in + let s_list = + List.concat + (List.rev + (List.map + ~f:(function + | _, x, _ -> x) + dse_list)) + in + let es = + List.map + ~f:(function + | _, _, x -> x) + dse_list + in + slist_concat_no_loc + (d_list @ s_list) + (match Map.find fim s with + | None -> NRFunApp (t, s, es) + | Some (_, args, b) -> + let b = replace_fresh_local_vars b in + let b = handle_early_returns None b in + (subst_args_stmt args es { pattern = b; meta = Location_span.empty }) + .pattern) + | Return e -> + (match e with | None -> Return None | Some e -> - let d, s, e = inline_function_expression adt fim e in - slist_concat_no_loc (d @ s) (Return (Some e)) ) - | IfElse (e, s1, s2) -> let d, s, e = inline_function_expression adt fim e in - slist_concat_no_loc (d @ s) - (IfElse - ( e - , inline_function_statement adt fim s1 - , Option.map ~f:(inline_function_statement adt fim) s2 )) + slist_concat_no_loc (d @ s) (Return (Some e))) + | IfElse (e, s1, s2) -> + let d, s, e = inline_function_expression adt fim e in + slist_concat_no_loc + (d @ s) + (IfElse + ( e + , inline_function_statement adt fim s1 + , Option.map ~f:(inline_function_statement adt fim) s2 )) | While (e, s) -> - let d', s', e = inline_function_expression adt fim e in - slist_concat_no_loc (d' @ s') - (While - ( e - , match s' with - | [] -> inline_function_statement adt fim s + let d', s', e = inline_function_expression adt fim e in + slist_concat_no_loc + (d' @ s') + (While + ( e + , match s' with + | [] -> inline_function_statement adt fim s + | _ -> + { pattern = + Block ([ inline_function_statement adt fim s ] @ map_no_loc s') + ; meta = Location_span.empty + } )) + | For { loopvar; lower; upper; body } -> + let d_lower, s_lower, lower = inline_function_expression adt fim lower in + let d_upper, s_upper, upper = inline_function_expression adt fim upper in + slist_concat_no_loc + (d_lower @ d_upper @ s_lower @ s_upper) + (For + { loopvar + ; lower + ; upper + ; body = + (match s_upper with + | [] -> inline_function_statement adt fim body | _ -> - { pattern= - Block - ( [inline_function_statement adt fim s] - @ map_no_loc s' ) - ; meta= Location_span.empty } )) - | For {loopvar; lower; upper; body} -> - let d_lower, s_lower, lower = - inline_function_expression adt fim lower - in - let d_upper, s_upper, upper = - inline_function_expression adt fim upper - in - slist_concat_no_loc - (d_lower @ d_upper @ s_lower @ s_upper) - (For - { loopvar - ; lower - ; upper - ; body= - ( match s_upper with - | [] -> inline_function_statement adt fim body - | _ -> - { pattern= - Block - ( [inline_function_statement adt fim body] - @ map_no_loc s_upper ) - ; meta= Location_span.empty } ) }) + { pattern = + Block + ([ inline_function_statement adt fim body ] + @ map_no_loc s_upper) + ; meta = Location_span.empty + }) + }) | Block l -> Block (List.map l ~f:(inline_function_statement adt fim)) | SList l -> SList (List.map l ~f:(inline_function_statement adt fim)) | Decl r -> Decl r | Skip -> Skip | Break -> Break - | Continue -> Continue ) - ; meta } + | Continue -> Continue) + ; meta + } +;; let create_function_inline_map adt l = (* We only add the first definition for each function to the inline map. @@ -432,27 +494,29 @@ let create_function_inline_map adt l = We also don't want to add any function declaration (as opposed to definitions), because that would replace the function call with a Skip. *) - let f (accum, visited) Program.({fdname; fdargs; fdbody; fdrt; _}) = - if Set.mem visited fdname then (accum, visited) - else + let f (accum, visited) Program.{ fdname; fdargs; fdbody; fdrt; _ } = + if Set.mem visited fdname + then accum, visited + else ( let accum' = match fdbody with - | Stmt.Fixed.({pattern= Stmt.Fixed.Pattern.Skip; _}) -> accum - | _ -> ( - let data = - ( Option.map ~f:(fun x -> Type.Unsized x) fdrt - , List.map ~f:(fun (_, name, _) -> name) fdargs - , inline_function_statement adt accum fdbody ) - in - match Map.add accum ~key:fdname ~data with - | `Ok m -> m - | `Duplicate -> accum ) + | Stmt.Fixed.{ pattern = Stmt.Fixed.Pattern.Skip; _ } -> accum + | _ -> + let data = + ( Option.map ~f:(fun x -> Type.Unsized x) fdrt + , List.map ~f:(fun (_, name, _) -> name) fdargs + , inline_function_statement adt accum fdbody ) + in + (match Map.add accum ~key:fdname ~data with + | `Ok m -> m + | `Duplicate -> accum) in let visited' = Set.add visited fdname in - (accum', visited') + accum', visited') in let accum, _ = List.fold l ~init:(Map.Poly.empty, Set.Poly.empty) ~f in accum +;; let function_inlining (mir : Program.Typed.t) = let dataonly_inline_map = @@ -462,138 +526,145 @@ let function_inlining (mir : Program.Typed.t) = create_function_inline_map UnsizedType.AutoDiffable mir.functions_block in let dataonly_inline_function_statements = - List.map - ~f:(inline_function_statement UnsizedType.DataOnly dataonly_inline_map) + List.map ~f:(inline_function_statement UnsizedType.DataOnly dataonly_inline_map) in let autodiffable_inline_function_statements = - List.map - ~f: - (inline_function_statement UnsizedType.AutoDiffable autodiff_inline_map) + List.map ~f:(inline_function_statement UnsizedType.AutoDiffable autodiff_inline_map) in { mir with - prepare_data= dataonly_inline_function_statements mir.prepare_data - ; transform_inits= - autodiffable_inline_function_statements mir.transform_inits - ; log_prob= autodiffable_inline_function_statements mir.log_prob - ; generate_quantities= - dataonly_inline_function_statements mir.generate_quantities } + prepare_data = dataonly_inline_function_statements mir.prepare_data + ; transform_inits = autodiffable_inline_function_statements mir.transform_inits + ; log_prob = autodiffable_inline_function_statements mir.log_prob + ; generate_quantities = dataonly_inline_function_statements mir.generate_quantities + } +;; -let rec contains_top_break_or_continue Stmt.Fixed.({pattern; _}) = +let rec contains_top_break_or_continue Stmt.Fixed.{ pattern; _ } = match pattern with | Break | Continue -> true | Assignment (_, _) - |TargetPE _ - |NRFunApp (_, _, _) - |Return _ | Decl _ - |While (_, _) - |For _ | Skip -> - false + | TargetPE _ + | NRFunApp (_, _, _) + | Return _ | Decl _ + | While (_, _) + | For _ | Skip -> false | Block l | SList l -> List.exists l ~f:contains_top_break_or_continue - | IfElse (_, b1, b2) -> ( - contains_top_break_or_continue b1 - || - match b2 with - | None -> false - | Some b -> contains_top_break_or_continue b ) + | IfElse (_, b1, b2) -> + contains_top_break_or_continue b1 + || + (match b2 with + | None -> false + | Some b -> contains_top_break_or_continue b) +;; let unroll_static_limit = 32 let unroll_static_loops_statement _ = let f stmt = match stmt with - | Stmt.Fixed.Pattern.For {loopvar; lower; upper; body} -> ( - let lower = Partial_evaluator.eval_expr lower in - let upper = Partial_evaluator.eval_expr upper in - match - (contains_top_break_or_continue body, lower.pattern, upper.pattern) - with - | false, Lit (Int, low_str), Lit (Int, up_str) -> - let low = Int.of_string low_str in - let up = Int.of_string up_str in - if up - low > unroll_static_limit then stmt - else - let range = - List.map - ~f:(fun i -> - Expr.Fixed. - { pattern= Lit (Int, Int.to_string i) - ; meta= - Expr.Typed.Meta. - { type_= UInt - ; loc= Location_span.empty - ; adlevel= DataOnly } } ) - (List.range ~start:`inclusive ~stop:`inclusive low up) - in - let stmts = - List.map - ~f:(fun i -> - subst_args_stmt [loopvar] [i] - {pattern= body.pattern; meta= Location_span.empty} ) - range - in - Stmt.Fixed.Pattern.SList stmts - | _ -> stmt ) + | Stmt.Fixed.Pattern.For { loopvar; lower; upper; body } -> + let lower = Partial_evaluator.eval_expr lower in + let upper = Partial_evaluator.eval_expr upper in + (match contains_top_break_or_continue body, lower.pattern, upper.pattern with + | false, Lit (Int, low_str), Lit (Int, up_str) -> + let low = Int.of_string low_str in + let up = Int.of_string up_str in + if up - low > unroll_static_limit + then stmt + else ( + let range = + List.map + ~f:(fun i -> + Expr.Fixed. + { pattern = Lit (Int, Int.to_string i) + ; meta = + Expr.Typed.Meta. + { type_ = UInt; loc = Location_span.empty; adlevel = DataOnly } + }) + (List.range ~start:`inclusive ~stop:`inclusive low up) + in + let stmts = + List.map + ~f:(fun i -> + subst_args_stmt + [ loopvar ] + [ i ] + { pattern = body.pattern; meta = Location_span.empty }) + range + in + Stmt.Fixed.Pattern.SList stmts) + | _ -> stmt) | _ -> stmt in top_down_map_rec_stmt_loc f +;; let static_loop_unrolling mir = transform_program_blockwise mir unroll_static_loops_statement +;; let unroll_loop_one_step_statement _ = let f stmt = match stmt with - | Stmt.Fixed.Pattern.For {loopvar; lower; upper; body} -> - if contains_top_break_or_continue body then stmt - else - IfElse - ( Expr.Fixed. - {lower with pattern= FunApp (StanLib, "Geq__", [upper; lower])} - , { pattern= - (let body_unrolled = - subst_args_stmt [loopvar] [lower] - {pattern= body.pattern; meta= Location_span.empty} - in - let (body' : Stmt.Located.t) = - { pattern= - Stmt.Fixed.Pattern.For - { loopvar - ; upper - ; body - ; lower= - { lower with - pattern= - FunApp - ( StanLib - , "Plus__" - , [lower; Expr.Helpers.loop_bottom] ) } } - ; meta= Location_span.empty } - in - match body_unrolled.pattern with - | Block stmts -> Block (stmts @ [body']) - | _ -> Stmt.Fixed.Pattern.Block [body_unrolled; body']) - ; meta= Location_span.empty } - , None ) + | Stmt.Fixed.Pattern.For { loopvar; lower; upper; body } -> + if contains_top_break_or_continue body + then stmt + else + IfElse + ( Expr.Fixed.{ lower with pattern = FunApp (StanLib, "Geq__", [ upper; lower ]) } + , { pattern = + (let body_unrolled = + subst_args_stmt + [ loopvar ] + [ lower ] + { pattern = body.pattern; meta = Location_span.empty } + in + let (body' : Stmt.Located.t) = + { pattern = + Stmt.Fixed.Pattern.For + { loopvar + ; upper + ; body + ; lower = + { lower with + pattern = + FunApp + (StanLib, "Plus__", [ lower; Expr.Helpers.loop_bottom ]) + } + } + ; meta = Location_span.empty + } + in + match body_unrolled.pattern with + | Block stmts -> Block (stmts @ [ body' ]) + | _ -> Stmt.Fixed.Pattern.Block [ body_unrolled; body' ]) + ; meta = Location_span.empty + } + , None ) | While (e, body) -> - if contains_top_break_or_continue body then stmt - else - IfElse - ( e - , { pattern= Block [body; {body with pattern= While (e, body)}] - ; meta= Location_span.empty } - , None ) + if contains_top_break_or_continue body + then stmt + else + IfElse + ( e + , { pattern = Block [ body; { body with pattern = While (e, body) } ] + ; meta = Location_span.empty + } + , None ) | _ -> stmt in map_rec_stmt_loc f +;; let one_step_loop_unrolling mir = transform_program_blockwise mir unroll_loop_one_step_statement +;; let collapse_lists_statement _ = let rec collapse_lists l = match l with | [] -> [] - | Stmt.Fixed.({pattern= SList l'; _}) :: rest -> l' @ collapse_lists rest + | Stmt.Fixed.{ pattern = SList l'; _ } :: rest -> l' @ collapse_lists rest | x :: rest -> x :: collapse_lists rest in let f = function @@ -602,77 +673,76 @@ let collapse_lists_statement _ = | x -> x in map_rec_stmt_loc f +;; let list_collapsing (mir : Program.Typed.t) = transform_program_blockwise mir collapse_lists_statement +;; let propagation (propagation_transfer : - (int, Stmt.Located.Non_recursive.t) Map.Poly.t - -> (module - Monotone_framework_sigs.TRANSFER_FUNCTION + (int, Stmt.Located.Non_recursive.t) Map.Poly.t + -> (module Monotone_framework_sigs.TRANSFER_FUNCTION with type labels = int - and type properties = (string, Middle.Expr.Typed.t) Map.Poly.t - option)) (mir : Program.Typed.t) = + and type properties = (string, Middle.Expr.Typed.t) Map.Poly.t option)) + (mir : Program.Typed.t) + = let transform s = - let flowgraph, flowgraph_to_mir = - Monotone_framework.forward_flowgraph_of_stmt s - in + let flowgraph, flowgraph_to_mir = Monotone_framework.forward_flowgraph_of_stmt s in let (module Flowgraph) = flowgraph in let values = - Monotone_framework.propagation_mfp mir + Monotone_framework.propagation_mfp + mir (module Flowgraph) - flowgraph_to_mir propagation_transfer + flowgraph_to_mir + propagation_transfer in let propagate_stmt = map_rec_stmt_loc_num flowgraph_to_mir (fun i -> subst_stmt_base - (Option.value ~default:Map.Poly.empty (Map.find_exn values i).entry) - ) + (Option.value ~default:Map.Poly.empty (Map.find_exn values i).entry)) in propagate_stmt (Map.find_exn flowgraph_to_mir 1) in transform_program mir transform +;; -let constant_propagation = - propagation Monotone_framework.constant_propagation_transfer +let constant_propagation = propagation Monotone_framework.constant_propagation_transfer let rec expr_any pred (e : Expr.Typed.t) = match e.pattern with | Indexed (e, is) -> expr_any pred e || List.exists ~f:(idx_any pred) is | _ -> pred e || Expr.Fixed.Pattern.fold (accum_any pred) false e.pattern -and idx_any pred (i : Expr.Typed.t Index.t) = - Index.fold (accum_any pred) false i - +and idx_any pred (i : Expr.Typed.t Index.t) = Index.fold (accum_any pred) false i and accum_any pred b e = b || expr_any pred e let can_side_effect_top_expr (e : Expr.Typed.t) = match e.pattern with | FunApp (t, f, _) -> - String.suffix f 3 = "_lp" - || (t = CompilerInternal && f = Internal_fun.to_string FnReadParam) - || (t = CompilerInternal && f = Internal_fun.to_string FnReadData) - || (t = CompilerInternal && f = Internal_fun.to_string FnWriteParam) - || (t = CompilerInternal && f = Internal_fun.to_string FnConstrain) - || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSize) - || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSize) - || t = CompilerInternal - && f = Internal_fun.to_string FnValidateSizeSimplex - || t = CompilerInternal - && f = Internal_fun.to_string FnValidateSizeUnitVector - || (t = CompilerInternal && f = Internal_fun.to_string FnUnconstrain) + String.suffix f 3 = "_lp" + || (t = CompilerInternal && f = Internal_fun.to_string FnReadParam) + || (t = CompilerInternal && f = Internal_fun.to_string FnReadData) + || (t = CompilerInternal && f = Internal_fun.to_string FnWriteParam) + || (t = CompilerInternal && f = Internal_fun.to_string FnConstrain) + || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSize) + || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSize) + || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSizeSimplex) + || (t = CompilerInternal && f = Internal_fun.to_string FnValidateSizeUnitVector) + || (t = CompilerInternal && f = Internal_fun.to_string FnUnconstrain) | _ -> false +;; let cannot_duplicate_expr (e : Expr.Typed.t) = let pred e = can_side_effect_top_expr e - || ( match e.pattern with + || (match e.pattern with | FunApp (_, f, _) -> String.suffix f 4 = "_rng" - | _ -> false ) + | _ -> false) || (preserve_stability && UnsizedType.is_autodiffable e.meta.type_) in expr_any pred e +;; let cannot_remove_expr (e : Expr.Typed.t) = expr_any can_side_effect_top_expr e @@ -680,32 +750,33 @@ let expression_propagation mir = propagation (Monotone_framework.expression_propagation_transfer cannot_duplicate_expr) mir +;; let copy_propagation mir = let globals = Monotone_framework.globals mir in propagation (Monotone_framework.copy_propagation_transfer globals) mir +;; let is_skip_break_continue s = match s with | Stmt.Fixed.Pattern.Skip | Break | Continue -> true | _ -> false +;; (* TODO: could also implement partial dead code elimination *) let dead_code_elimination (mir : Program.Typed.t) = (* TODO: think about whether we should treat function bodies as local scopes in the statement - from the POV of a live variables analysis. - (Obviously, this shouldn't be the case for the purposes of reaching definitions, - constant propagation, expressions analyses. But I do think that's the right way to - go about live variables. *) + from the POV of a live variables analysis. + (Obviously, this shouldn't be the case for the purposes of reaching definitions, + constant propagation, expressions analyses. But I do think that's the right way to + go about live variables. *) let transform s = let rev_flowgraph, flowgraph_to_mir = Monotone_framework.inverse_flowgraph_of_stmt s in let (module Rev_Flowgraph) = rev_flowgraph in let live_variables = - Monotone_framework.live_variables_mfp mir - (module Rev_Flowgraph) - flowgraph_to_mir + Monotone_framework.live_variables_mfp mir (module Rev_Flowgraph) flowgraph_to_mir in let dead_code_elim_stmt_base i stmt = (* NOTE: entry in the reverse flowgraph, so exit in the forward flowgraph *) @@ -714,59 +785,53 @@ let dead_code_elimination (mir : Program.Typed.t) = in match stmt with | Stmt.Fixed.Pattern.Assignment ((x, _, []), rhs) -> - if Set.Poly.mem live_variables_s x || cannot_remove_expr rhs then - stmt - else Skip + if Set.Poly.mem live_variables_s x || cannot_remove_expr rhs then stmt else Skip | Assignment ((x, _, is), rhs) -> - if - Set.Poly.mem live_variables_s x - || cannot_remove_expr rhs - || List.exists ~f:(idx_any cannot_remove_expr) is - then stmt - else Skip + if Set.Poly.mem live_variables_s x + || cannot_remove_expr rhs + || List.exists ~f:(idx_any cannot_remove_expr) is + then stmt + else Skip (* NOTE: we never get rid of declarations as we might not be able to - remove an assignment to a variable - due to side effects. *) + remove an assignment to a variable + due to side effects. *) (* TODO: maybe we should revisit that. *) - | Decl _ | TargetPE _ - |NRFunApp (_, _, _) - |Break | Continue | Return _ | Skip -> - stmt - | IfElse (e, b1, b2) -> ( - if - (* TODO: check if e has side effects, like print, reject, then don't optimize? *) - (not (cannot_remove_expr e)) - && b1.Stmt.Fixed.pattern = Skip - && ( Option.map ~f:(fun Stmt.Fixed.({pattern; _}) -> pattern) b2 - = Some Skip - || Option.map ~f:(fun Stmt.Fixed.({pattern; _}) -> pattern) b2 - = None ) - then Skip - else - match e.pattern with - | Lit (Int, "0") | Lit (Real, "0.0") -> ( - match b2 with Some x -> x.pattern | None -> Skip ) - | Lit (_, _) -> b1.pattern - | _ -> IfElse (e, b1, b2) ) - | While (e, b) -> ( - if (not (cannot_remove_expr e)) && b.pattern = Break then Skip - else - match e.pattern with - | Lit (Int, "0") | Lit (Real, "0.0") -> Skip - | _ -> While (e, b) ) - | For {loopvar; lower; upper; body} -> - if - (not (cannot_remove_expr lower)) - && (not (cannot_remove_expr upper)) - && is_skip_break_continue body.pattern - then Skip - else For {loopvar; lower; upper; body} + | Decl _ | TargetPE _ | NRFunApp (_, _, _) | Break | Continue | Return _ | Skip -> + stmt + | IfElse (e, b1, b2) -> + if (* TODO: check if e has side effects, like print, reject, then don't optimize? *) + (not (cannot_remove_expr e)) + && b1.Stmt.Fixed.pattern = Skip + && (Option.map ~f:(fun Stmt.Fixed.{ pattern; _ } -> pattern) b2 = Some Skip + || Option.map ~f:(fun Stmt.Fixed.{ pattern; _ } -> pattern) b2 = None) + then Skip + else ( + match e.pattern with + | Lit (Int, "0") | Lit (Real, "0.0") -> + (match b2 with + | Some x -> x.pattern + | None -> Skip) + | Lit (_, _) -> b1.pattern + | _ -> IfElse (e, b1, b2)) + | While (e, b) -> + if (not (cannot_remove_expr e)) && b.pattern = Break + then Skip + else ( + match e.pattern with + | Lit (Int, "0") | Lit (Real, "0.0") -> Skip + | _ -> While (e, b)) + | For { loopvar; lower; upper; body } -> + if (not (cannot_remove_expr lower)) + && (not (cannot_remove_expr upper)) + && is_skip_break_continue body.pattern + then Skip + else For { loopvar; lower; upper; body } | Block l -> - let l' = List.filter ~f:(fun x -> x.pattern <> Skip) l in - if List.length l' = 0 then Skip else Block l' + let l' = List.filter ~f:(fun x -> x.pattern <> Skip) l in + if List.length l' = 0 then Skip else Block l' | SList l -> - let l' = List.filter ~f:(fun x -> x.pattern <> Skip) l in - SList l' + let l' = List.filter ~f:(fun x -> x.pattern <> Skip) l in + SList l' in let dead_code_elim_stmt = map_rec_stmt_loc_num flowgraph_to_mir dead_code_elim_stmt_base @@ -774,6 +839,7 @@ let dead_code_elimination (mir : Program.Typed.t) = dead_code_elim_stmt (Map.find_exn flowgraph_to_mir 1) in transform_program mir transform +;; let partial_evaluation = Partial_evaluator.eval_prog @@ -783,38 +849,43 @@ let lazy_code_motion (mir : Program.Typed.t) = simultaneously *) let preprocess_flowgraph = let preprocess_flowgraph_base - (stmt : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) = + (stmt : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) + = match stmt with | IfElse (e, b1, Some b2) -> - Stmt.Fixed.( - Pattern.IfElse - ( e - , { pattern= - Block [b1; {pattern= Skip; meta= Location_span.empty}] - ; meta= Location_span.empty } - , Some - { pattern= - Block [b2; {pattern= Skip; meta= Location_span.empty}] - ; meta= Location_span.empty } )) - | IfElse (e, b, None) -> - IfElse + Stmt.Fixed.( + Pattern.IfElse ( e - , { pattern= Block [b; {pattern= Skip; meta= Location_span.empty}] - ; meta= Location_span.empty } - , Some {pattern= Skip; meta= Location_span.empty} ) + , { pattern = Block [ b1; { pattern = Skip; meta = Location_span.empty } ] + ; meta = Location_span.empty + } + , Some + { pattern = Block [ b2; { pattern = Skip; meta = Location_span.empty } ] + ; meta = Location_span.empty + } )) + | IfElse (e, b, None) -> + IfElse + ( e + , { pattern = Block [ b; { pattern = Skip; meta = Location_span.empty } ] + ; meta = Location_span.empty + } + , Some { pattern = Skip; meta = Location_span.empty } ) | While (e, b) -> - While - ( e - , { pattern= Block [b; {pattern= Skip; meta= Location_span.empty}] - ; meta= Location_span.empty } ) - | For {loopvar; lower; upper; body= b} -> - For - { loopvar - ; lower - ; upper - ; body= - { pattern= Block [b; {pattern= Skip; meta= Location_span.empty}] - ; meta= Location_span.empty } } + While + ( e + , { pattern = Block [ b; { pattern = Skip; meta = Location_span.empty } ] + ; meta = Location_span.empty + } ) + | For { loopvar; lower; upper; body = b } -> + For + { loopvar + ; lower + ; upper + ; body = + { pattern = Block [ b; { pattern = Skip; meta = Location_span.empty } ] + ; meta = Location_span.empty + } + } | _ -> stmt in map_rec_stmt_loc preprocess_flowgraph_base @@ -825,8 +896,7 @@ let lazy_code_motion (mir : Program.Typed.t) = in let fwd_flowgraph = Monotone_framework.reverse rev_flowgraph in let latest_expr, used_not_latest_expressions_mfp = - Monotone_framework.lazy_expressions_mfp fwd_flowgraph rev_flowgraph - flowgraph_to_mir + Monotone_framework.lazy_expressions_mfp fwd_flowgraph rev_flowgraph flowgraph_to_mir in let expression_map = let rec collect_expressions accum (e : Expr.Typed.t) = @@ -834,25 +904,28 @@ let lazy_code_motion (mir : Program.Typed.t) = | Lit (_, _) -> accum | Var _ -> accum | _ when cannot_duplicate_expr e -> - (* Immovable expressions might have movable subexpressions *) - Expr.Fixed.Pattern.fold collect_expressions accum e.pattern + (* Immovable expressions might have movable subexpressions *) + Expr.Fixed.Pattern.fold collect_expressions accum e.pattern | _ -> Map.set accum ~key:e ~data:(Gensym.generate ~prefix:"lcm_" ()) in Set.fold (Monotone_framework.used_expressions_stmt s.pattern) - ~init:Expr.Typed.Map.empty ~f:collect_expressions + ~init:Expr.Typed.Map.empty + ~f:collect_expressions in (* TODO: it'd be more efficient to just not accumulate constants in the static analysis *) let declarations_list = Map.fold expression_map ~init:[] ~f:(fun ~key ~data accum -> Stmt.Fixed. - { pattern= + { pattern = Pattern.Decl - { decl_adtype= Expr.Typed.adlevel_of key - ; decl_id= data - ; decl_type= Type.Unsized (Expr.Typed.type_of key) } - ; meta= Location_span.empty } - :: accum ) + { decl_adtype = Expr.Typed.adlevel_of key + ; decl_id = data + ; decl_type = Type.Unsized (Expr.Typed.type_of key) + } + ; meta = Location_span.empty + } + :: accum) in let lazy_code_motion_base i stmt = let latest_and_used_after_i = @@ -864,8 +937,7 @@ let lazy_code_motion (mir : Program.Typed.t) = latest_and_used_after_i |> Set.filter ~f:(fun x -> Map.mem expression_map x) |> Set.to_list - |> List.sort ~compare:(fun e e' -> - compare_int (expr_depth e) (expr_depth e') ) + |> List.sort ~compare:(fun e e' -> compare_int (expr_depth e) (expr_depth e')) in (* TODO: is this sort doing anything or are they already stored in the right order by chance? It appears to not do anything. *) @@ -873,10 +945,9 @@ let lazy_code_motion (mir : Program.Typed.t) = List.map ~f:(fun e -> Stmt.Fixed. - { pattern= - Assignment - ((Map.find_exn expression_map e, e.meta.type_, []), e) - ; meta= Location_span.empty } ) + { pattern = Assignment ((Map.find_exn expression_map e, e.meta.type_, []), e) + ; meta = Location_span.empty + }) to_assign_in_s in let expr_subst_stmt_except_initial_assign m = @@ -884,9 +955,8 @@ let lazy_code_motion (mir : Program.Typed.t) = match stmt with | Stmt.Fixed.Pattern.Assignment ((x, _, []), e') when Map.mem m e' - && Expr.Typed.equal {e' with pattern= Var x} - (Map.find_exn m e') -> - expr_subst_stmt_base (Map.remove m e') stmt + && Expr.Typed.equal { e' with pattern = Var x } (Map.find_exn m e') -> + expr_subst_stmt_base (Map.remove m e') stmt | _ -> expr_subst_stmt_base m stmt in map_rec_stmt_loc f @@ -895,81 +965,79 @@ let lazy_code_motion (mir : Program.Typed.t) = Map.filter_keys ~f:(fun key -> Set.mem latest_and_used_after_i key - || Set.mem (Map.find_exn used_not_latest_expressions_mfp i).exit - key ) - (Map.mapi expression_map ~f:(fun ~key ~data -> - {key with pattern= Var data} )) + || Set.mem (Map.find_exn used_not_latest_expressions_mfp i).exit key) + (Map.mapi expression_map ~f:(fun ~key ~data -> { key with pattern = Var data })) in let f = expr_subst_stmt_except_initial_assign expr_map in - if List.length assignments_to_add_to_s = 0 then - (f Stmt.Fixed.{pattern= stmt; meta= Location_span.empty}).pattern + if List.length assignments_to_add_to_s = 0 + then (f Stmt.Fixed.{ pattern = stmt; meta = Location_span.empty }).pattern else SList - (List.map ~f - ( assignments_to_add_to_s - @ [{pattern= stmt; meta= Location_span.empty}] )) + (List.map + ~f + (assignments_to_add_to_s @ [ { pattern = stmt; meta = Location_span.empty } ])) in let lazy_code_motion_stmt = map_rec_stmt_loc_num flowgraph_to_mir lazy_code_motion_base in Stmt.Fixed. - { pattern= + { pattern = SList - ( declarations_list - @ [lazy_code_motion_stmt (Map.find_exn flowgraph_to_mir 1)] ) - ; meta= Location_span.empty } + (declarations_list + @ [ lazy_code_motion_stmt (Map.find_exn flowgraph_to_mir 1) ]) + ; meta = Location_span.empty + } in let cleanup = - let cleanup_base - (stmt : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) : - (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t = + let cleanup_base (stmt : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t) + : (Expr.Typed.t, Stmt.Located.t) Stmt.Fixed.Pattern.t + = match stmt with - | Stmt.Fixed.(Pattern.IfElse - ( e - , {pattern= Block [b1; {pattern= Skip; _}]; _} - , Some {pattern= Block [b2; {pattern= Skip; _}]; _} )) -> - IfElse (e, b1, Some b2) + | Stmt.Fixed.( + Pattern.IfElse + ( e + , { pattern = Block [ b1; { pattern = Skip; _ } ]; _ } + , Some { pattern = Block [ b2; { pattern = Skip; _ } ]; _ } )) -> + IfElse (e, b1, Some b2) | IfElse ( e - , {pattern= Block [b; {pattern= Skip; _}]; _} - , Some {pattern= Skip; _} ) -> - IfElse (e, b, None) - | While (e, {pattern= Block [b; {pattern= Skip; _}]; _}) -> While (e, b) + , { pattern = Block [ b; { pattern = Skip; _ } ]; _ } + , Some { pattern = Skip; _ } ) -> IfElse (e, b, None) + | While (e, { pattern = Block [ b; { pattern = Skip; _ } ]; _ }) -> While (e, b) | For { loopvar ; lower ; upper - ; body= {pattern= Block [b; {pattern= Skip; _}]; _} } -> - For {loopvar; lower; upper; body= b} + ; body = { pattern = Block [ b; { pattern = Skip; _ } ]; _ } + } -> For { loopvar; lower; upper; body = b } | _ -> stmt in map_rec_stmt_loc cleanup_base in transform_program_blockwise mir (fun _ x -> - cleanup (transform (preprocess_flowgraph x)) ) + cleanup (transform (preprocess_flowgraph x))) +;; let block_fixing mir = transform_program_blockwise mir (fun _ x -> (map_rec_stmt_loc (fun stmt -> match stmt with | IfElse - ( e - , {pattern= SList l; meta} - , Some {pattern= SList l'; meta= smeta'} ) -> - IfElse - ( e - , {pattern= Block l; meta} - , Some {pattern= Block l'; meta= smeta'} ) - | IfElse (e, {pattern= SList l; meta}, b) -> - IfElse (e, {pattern= Block l; meta}, b) - | IfElse (e, b, Some {pattern= SList l'; meta= smeta'}) -> - IfElse (e, b, Some {pattern= Block l'; meta= smeta'}) - | While (e, {pattern= SList l; meta}) -> - While (e, {pattern= Block l; meta}) - | For {loopvar; lower; upper; body= {pattern= SList l; meta}} -> - For {loopvar; lower; upper; body= {pattern= Block l; meta}} - | _ -> stmt )) - x ) + (e, { pattern = SList l; meta }, Some { pattern = SList l'; meta = smeta' }) + -> + IfElse + (e, { pattern = Block l; meta }, Some { pattern = Block l'; meta = smeta' }) + | IfElse (e, { pattern = SList l; meta }, b) -> + IfElse (e, { pattern = Block l; meta }, b) + | IfElse (e, b, Some { pattern = SList l'; meta = smeta' }) -> + IfElse (e, b, Some { pattern = Block l'; meta = smeta' }) + | While (e, { pattern = SList l; meta }) -> + While (e, { pattern = Block l; meta }) + | For { loopvar; lower; upper; body = { pattern = SList l; meta } } -> + For { loopvar; lower; upper; body = { pattern = Block l; meta } } + | _ -> stmt)) + x) +;; (* TODO: implement SlicStan style optimizer for choosing best program block for each statement. *) (* TODO: add optimization pass to move declarations down as much as possible and introduce as @@ -981,8 +1049,10 @@ let optimize_ad_levels (mir : Program.Typed.t) = let global_initial_ad_variables = Set.Poly.of_list (List.filter_map - ~f:(fun (v, Program.({out_block; _})) -> - match out_block with Parameters -> Some v | _ -> None ) + ~f:(fun (v, Program.{ out_block; _ }) -> + match out_block with + | Parameters -> Some v + | _ -> None) mir.output_vars) in let transform fundef_opt s = @@ -995,18 +1065,19 @@ let optimize_ad_levels (mir : Program.Typed.t) = let initial_ad_variables = match (fundef_opt : Stmt.Located.t Program.fun_def option) with | None -> global_initial_ad_variables - | Some {fdargs; _} -> - Set.Poly.union global_initial_ad_variables - (Set.Poly.of_list - (List.filter_map fdargs ~f:(fun (_, name, ut) -> - if UnsizedType.is_autodiffable ut then Some name else None - ))) + | Some { fdargs; _ } -> + Set.Poly.union + global_initial_ad_variables + (Set.Poly.of_list + (List.filter_map fdargs ~f:(fun (_, name, ut) -> + if UnsizedType.is_autodiffable ut then Some name else None))) in let ad_levels = Monotone_framework.autodiff_level_mfp (module Fwd_Flowgraph) (module Rev_Flowgraph) - flowgraph_to_mir initial_ad_variables + flowgraph_to_mir + initial_ad_variables in let insert_constraint_variables vars = Set.Poly.union vars (Set.Poly.map ~f:(fun x -> x ^ "_in__") vars) @@ -1021,12 +1092,10 @@ let optimize_ad_levels (mir : Program.Typed.t) = (fun x -> x) stmt with - | Decl {decl_id; decl_type; _} - when Set.mem autodiffable_variables decl_id -> - Stmt.Fixed.Pattern.Decl - {decl_adtype= AutoDiffable; decl_id; decl_type} - | Decl {decl_id; decl_type; _} -> - Decl {decl_adtype= DataOnly; decl_id; decl_type} + | Decl { decl_id; decl_type; _ } when Set.mem autodiffable_variables decl_id -> + Stmt.Fixed.Pattern.Decl { decl_adtype = AutoDiffable; decl_id; decl_type } + | Decl { decl_id; decl_type; _ } -> + Decl { decl_adtype = DataOnly; decl_id; decl_type } | s -> s in let optimize_ad_levels_stmt = @@ -1035,36 +1104,40 @@ let optimize_ad_levels (mir : Program.Typed.t) = optimize_ad_levels_stmt (Map.find_exn flowgraph_to_mir 1) in transform_program_blockwise mir transform +;; (* Apparently you need to completely copy/paste type definitions between ml and mli files?*) type optimization_settings = - { function_inlining: bool - ; static_loop_unrolling: bool - ; one_step_loop_unrolling: bool - ; list_collapsing: bool - ; block_fixing: bool - ; constant_propagation: bool - ; expression_propagation: bool - ; copy_propagation: bool - ; dead_code_elimination: bool - ; partial_evaluation: bool - ; lazy_code_motion: bool - ; optimize_ad_levels: bool } + { function_inlining : bool + ; static_loop_unrolling : bool + ; one_step_loop_unrolling : bool + ; list_collapsing : bool + ; block_fixing : bool + ; constant_propagation : bool + ; expression_propagation : bool + ; copy_propagation : bool + ; dead_code_elimination : bool + ; partial_evaluation : bool + ; lazy_code_motion : bool + ; optimize_ad_levels : bool + } let settings_const b = - { function_inlining= b - ; static_loop_unrolling= b - ; one_step_loop_unrolling= b - ; list_collapsing= b - ; block_fixing= b - ; constant_propagation= b - ; expression_propagation= b - ; copy_propagation= b - ; dead_code_elimination= b - ; partial_evaluation= b - ; lazy_code_motion= b - ; optimize_ad_levels= b } + { function_inlining = b + ; static_loop_unrolling = b + ; one_step_loop_unrolling = b + ; list_collapsing = b + ; block_fixing = b + ; constant_propagation = b + ; expression_propagation = b + ; copy_propagation = b + ; dead_code_elimination = b + ; partial_evaluation = b + ; lazy_code_motion = b + ; optimize_ad_levels = b + } +;; let all_optimizations : optimization_settings = settings_const true let no_optimizations : optimization_settings = settings_const false @@ -1075,43 +1148,44 @@ let optimization_suite ?(settings = all_optimizations) mir = (* Book section A *) (* Book section B *) (* Book: Procedure integration *) - (function_inlining, settings.function_inlining) + function_inlining, settings.function_inlining (* Book: Sparse conditional constant propagation *) - ; (constant_propagation, settings.constant_propagation) + ; constant_propagation, settings.constant_propagation (* Book section C *) (* Book: Local and global copy propagation *) - ; (copy_propagation, settings.copy_propagation) + ; copy_propagation, settings.copy_propagation (* Book: Sparse conditional constant propagation *) - ; (constant_propagation, settings.constant_propagation) + ; constant_propagation, settings.constant_propagation (* Book: Dead-code elimination *) - ; (dead_code_elimination, settings.dead_code_elimination) + ; dead_code_elimination, settings.dead_code_elimination (* Matthijs: Before lazy code motion to get loop-invariant code motion *) - ; (one_step_loop_unrolling, settings.one_step_loop_unrolling) + ; one_step_loop_unrolling, settings.one_step_loop_unrolling (* Matthjis: expression_propagation < partial_evaluation *) - ; (expression_propagation, settings.expression_propagation) + ; expression_propagation, settings.expression_propagation (* Matthjis: partial_evaluation < lazy_code_motion *) - ; (partial_evaluation, settings.partial_evaluation) + ; partial_evaluation, settings.partial_evaluation (* Book: Loop-invariant code motion *) - ; (lazy_code_motion, settings.lazy_code_motion) + ; lazy_code_motion, settings.lazy_code_motion (* Matthijs: lazy_code_motion < copy_propagation TODO: Check if this is necessary *) - ; (copy_propagation, settings.copy_propagation) + ; copy_propagation, settings.copy_propagation (* Matthijs: Constant propagation before static loop unrolling *) - ; (constant_propagation, settings.constant_propagation) - (* Book: Loop simplification *) - ; (static_loop_unrolling, settings.static_loop_unrolling) + ; constant_propagation, settings.constant_propagation (* Book: Loop simplification *) + ; static_loop_unrolling, settings.static_loop_unrolling (* Book: Dead-code elimination *) (* Matthijs: Everything < Dead-code elimination *) - ; (dead_code_elimination, settings.dead_code_elimination) + ; dead_code_elimination, settings.dead_code_elimination (* Book: Machine idioms and instruction combining *) - ; (list_collapsing, settings.list_collapsing) + ; list_collapsing, settings.list_collapsing (* Book: Machine idioms and instruction combining *) - ; (optimize_ad_levels, settings.optimize_ad_levels) + ; optimize_ad_levels, settings.optimize_ad_levels (* Book: Machine idioms and instruction combining *) (* Matthijs: Everything < block_fixing *) - ; (block_fixing, settings.block_fixing) ] + ; block_fixing, settings.block_fixing + ] in let optimizations = List.filter_map maybe_optimizations ~f:(fun (fn, flag) -> - if flag then Some fn else None ) + if flag then Some fn else None) in List.fold optimizations ~init:mir ~f:(fun mir opt -> opt mir) +;; diff --git a/src/analysis_and_optimization/Optimize.mli b/src/analysis_and_optimization/Optimize.mli index aee939f47b..75e51050f9 100644 --- a/src/analysis_and_optimization/Optimize.mli +++ b/src/analysis_and_optimization/Optimize.mli @@ -1,78 +1,81 @@ (* Code for optimization passes on the MIR *) open Middle -val function_inlining : Program.Typed.t -> Program.Typed.t (** Inline all functions except for ones with forward declarations (e.g. recursive functions, mutually recursive functions, and functions without a definition *) +val function_inlining : Program.Typed.t -> Program.Typed.t -val static_loop_unrolling : Program.Typed.t -> Program.Typed.t (** Unroll all for-loops with constant bounds, as long as they do not contain break or continue statements in their body at the top level *) +val static_loop_unrolling : Program.Typed.t -> Program.Typed.t -val one_step_loop_unrolling : Program.Typed.t -> Program.Typed.t (** Unroll all loops for one iteration, as long as they do not contain break or continue statements in their body at the top level *) +val one_step_loop_unrolling : Program.Typed.t -> Program.Typed.t -val list_collapsing : Program.Typed.t -> Program.Typed.t (** Remove redundant SList constructors from the Mir that might have been introduced by other optimizations *) +val list_collapsing : Program.Typed.t -> Program.Typed.t -val block_fixing : Program.Typed.t -> Program.Typed.t (** Make sure that SList constructors directly under if, for, while or fundef constructors are replaced with Block constructors. This should probably be run before we generate code. *) +val block_fixing : Program.Typed.t -> Program.Typed.t -val constant_propagation : Program.Typed.t -> Program.Typed.t (** Propagate constant values through variable assignments *) +val constant_propagation : Program.Typed.t -> Program.Typed.t -val expression_propagation : Program.Typed.t -> Program.Typed.t (** Propagate arbitrary expressions through variable assignments. This can be useful for opening up new possibilities for partial evaluation. It should be followed by some CSE or lazy code motion pass, however. *) +val expression_propagation : Program.Typed.t -> Program.Typed.t -val copy_propagation : Program.Typed.t -> Program.Typed.t (** Propagate copies of variables through assignments. *) +val copy_propagation : Program.Typed.t -> Program.Typed.t -val dead_code_elimination : Program.Typed.t -> Program.Typed.t (** Eliminate semantically redundant code branches. This includes removing redundant assignments (because they will be overwritten) and removing redundant code in program branches that will never be reached. *) +val dead_code_elimination : Program.Typed.t -> Program.Typed.t -val partial_evaluation : Program.Typed.t -> Program.Typed.t (** Partially evaluate expressions in the program. This includes simplification using algebraic identities of logical and arithmetic operators as well as Stan math functions. *) +val partial_evaluation : Program.Typed.t -> Program.Typed.t -val lazy_code_motion : Program.Typed.t -> Program.Typed.t (** Perform partial redundancy elmination using the lazy code motion algorithm. This subsumes common subexpression elimination and loop-invariant code motion. *) +val lazy_code_motion : Program.Typed.t -> Program.Typed.t -val optimize_ad_levels : Program.Typed.t -> Program.Typed.t (** Assign the optimal ad-levels to local variables. That means, make sure that variables only ever get treated as autodiff variables if they have some dependency on a parameter *) +val optimize_ad_levels : Program.Typed.t -> Program.Typed.t (** Interface for turning individual optimizations on/off. Useful for testing and for top-level interface flags. *) type optimization_settings = - { function_inlining: bool - ; static_loop_unrolling: bool - ; one_step_loop_unrolling: bool - ; list_collapsing: bool - ; block_fixing: bool - ; constant_propagation: bool - ; expression_propagation: bool - ; copy_propagation: bool - ; dead_code_elimination: bool - ; partial_evaluation: bool - ; lazy_code_motion: bool - ; optimize_ad_levels: bool } + { function_inlining : bool + ; static_loop_unrolling : bool + ; one_step_loop_unrolling : bool + ; list_collapsing : bool + ; block_fixing : bool + ; constant_propagation : bool + ; expression_propagation : bool + ; copy_propagation : bool + ; dead_code_elimination : bool + ; partial_evaluation : bool + ; lazy_code_motion : bool + ; optimize_ad_levels : bool + } val all_optimizations : optimization_settings val no_optimizations : optimization_settings -val optimization_suite : - ?settings:optimization_settings -> Program.Typed.t -> Program.Typed.t (** Perform all optimizations in this module on the MIR in an appropriate order. *) +val optimization_suite + : ?settings:optimization_settings + -> Program.Typed.t + -> Program.Typed.t diff --git a/src/analysis_and_optimization/Partial_evaluator.ml b/src/analysis_and_optimization/Partial_evaluator.ml index c33d72b639..04ccc77b55 100644 --- a/src/analysis_and_optimization/Partial_evaluator.ml +++ b/src/analysis_and_optimization/Partial_evaluator.ml @@ -6,37 +6,39 @@ open Middle let preserve_stability = false -let is_int i Expr.Fixed.({pattern; _}) = - let nums = List.map ~f:(fun s -> string_of_int i ^ s) [""; "."; ".0"] in +let is_int i Expr.Fixed.{ pattern; _ } = + let nums = List.map ~f:(fun s -> string_of_int i ^ s) [ ""; "."; ".0" ] in match pattern with - | (Lit (Int, i) | Lit (Real, i)) when List.mem nums i ~equal:String.equal -> - true + | (Lit (Int, i) | Lit (Real, i)) when List.mem nums i ~equal:String.equal -> true | _ -> false +;; let apply_prefix_operator_int (op : string) i = Expr.Fixed.Pattern.Lit ( Int , Int.to_string - ( match op with + (match op with | "PPlus__" -> i | "PMinus__" -> -i | "PNot__" -> if i = 0 then 1 else 0 - | s -> raise_s [%sexp (s : string)] ) ) + | s -> raise_s [%sexp (s : string)]) ) +;; let apply_prefix_operator_real (op : string) i = Expr.Fixed.Pattern.Lit ( Real , Float.to_string - ( match op with + (match op with | "PPlus__" -> i | "PMinus__" -> -.i - | s -> raise_s [%sexp (s : string)] ) ) + | s -> raise_s [%sexp (s : string)]) ) +;; let apply_operator_int (op : string) i1 i2 = Expr.Fixed.Pattern.Lit ( Int , Int.to_string - ( match op with + (match op with | "Plus__" -> i1 + i2 | "Minus__" -> i1 - i2 | "Times__" -> i1 * i2 @@ -49,730 +51,803 @@ let apply_operator_int (op : string) i1 i2 = | "Leq__" -> Bool.to_int (i1 <= i2) | "Greater__" -> Bool.to_int (i1 > i2) | "Geq__" -> Bool.to_int (i1 >= i2) - | s -> raise_s [%sexp (s : string)] ) ) + | s -> raise_s [%sexp (s : string)]) ) +;; let apply_arithmetic_operator_real (op : string) r1 r2 = Expr.Fixed.Pattern.Lit ( Real , Float.to_string - ( match op with + (match op with | "Plus__" -> r1 +. r2 | "Minus__" -> r1 -. r2 | "Times__" -> r1 *. r2 | "Divide__" -> r1 /. r2 - | s -> raise_s [%sexp (s : string)] ) ) + | s -> raise_s [%sexp (s : string)]) ) +;; let apply_logical_operator_real (op : string) r1 r2 = Expr.Fixed.Pattern.Lit ( Int , Int.to_string - ( match op with + (match op with | "Equals__" -> Bool.to_int (r1 = r2) | "NEquals__" -> Bool.to_int (r1 <> r2) | "Less__" -> Bool.to_int (r1 < r2) | "Leq__" -> Bool.to_int (r1 <= r2) | "Greater__" -> Bool.to_int (r1 > r2) | "Geq__" -> Bool.to_int (r1 >= r2) - | s -> raise_s [%sexp (s : string)] ) ) + | s -> raise_s [%sexp (s : string)]) ) +;; let is_multi_index = function | Index.MultiIndex _ | Upfrom _ | Between _ | All -> true | Single _ -> false +;; let rec eval_expr (e : Expr.Typed.t) = { e with - pattern= - ( match e.pattern with + pattern = + (match e.pattern with | Var _ | Lit (_, _) -> e.pattern | FunApp (t, f, l) -> - let l = List.map ~f:eval_expr l in - let get_fun_or_op_rt_opt name l' = - let argument_types = - List.map ~f:(fun x -> Expr.Typed.(adlevel_of x, type_of x)) l' - in - Operator.of_string_opt name - |> Option.value_map - ~f:(fun op -> - Stan_math_signatures.operator_stan_math_return_type op - argument_types ) - ~default: - (Stan_math_signatures.stan_math_returntype name - argument_types) + let l = List.map ~f:eval_expr l in + let get_fun_or_op_rt_opt name l' = + let argument_types = + List.map ~f:(fun x -> Expr.Typed.(adlevel_of x, type_of x)) l' in - let try_partially_evaluate_to e = - Expr.Fixed.Pattern.( - match e with - | FunApp (StanLib, f', l') -> ( - match get_fun_or_op_rt_opt f' l' with - | Some _ -> FunApp (StanLib, f', l') - | None -> FunApp (StanLib, f, l) ) - | e -> e) - in - try_partially_evaluate_to - ( match (f, l) with - (* TODO: deal with tilde statements and unnormalized distributions properly here *) - | ( "bernoulli_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "inv_logit" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } ] ); _ } ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "bernoulli_logit_glm_lpmf", [y; x; alpha; beta]) - | ( "bernoulli_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "inv_logit" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } - ; alpha ] ); _ } ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "bernoulli_logit_glm_lpmf", [y; x; alpha; beta]) - | ( "bernoulli_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "inv_logit" - , [{pattern= FunApp (StanLib, "Times__", [x; beta]); _}] - ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "bernoulli_logit_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta] ) - | ( "bernoulli_logit_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "bernoulli_logit_glm_lpmf", [y; x; alpha; beta]) - | ( "bernoulli_logit_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ; alpha ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "bernoulli_logit_glm_lpmf", [y; x; alpha; beta]) - | ( "bernoulli_logit_lpmf" - , [y; {pattern= FunApp (StanLib, "Times__", [x; beta]); _}] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "bernoulli_logit_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta] ) - | ( "bernoulli_lpmf" - , [y; {pattern= FunApp (StanLib, "inv_logit", [alpha]); _}] ) -> - FunApp (StanLib, "bernoulli_logit_lpmf", [y; alpha]) - | ( "bernoulli_rng" - , [{pattern= FunApp (StanLib, "inv_logit", [alpha]); _}] ) -> - FunApp (StanLib, "bernoulli_logit_rng", [alpha]) - | ( "binomial_lpmf" - , [y; n; {pattern= FunApp (StanLib, "inv_logit", [alpha]); _}] ) - -> - FunApp (StanLib, "binomial_logit_lpmf", [y; n; alpha]) - | ( "categorical_lpmf" - , [y; {pattern= FunApp (StanLib, "inv_logit", [alpha]); _}] ) -> - FunApp (StanLib, "categorical_logit_lpmf", [y; alpha]) - | ( "categorical_rng" - , [{pattern= FunApp (StanLib, "inv_logit", [alpha]); _}] ) -> - FunApp (StanLib, "categorical_logit_rng", [alpha]) - | "columns_dot_product", [x; y] when Expr.Typed.equal x y -> - FunApp (StanLib, "columns_dot_self", [x]) - | "dot_product", [x; y] when Expr.Typed.equal x y -> - FunApp (StanLib, "dot_self", [x]) - | "inv", [{pattern= FunApp (StanLib, "sqrt", l); _}] -> - FunApp (StanLib, "inv_sqrt", l) - | "inv", [{pattern= FunApp (StanLib, "square", [x]); _}] -> - FunApp (StanLib, "inv_square", [x]) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "Minus__" - , [y; {pattern= FunApp (StanLib, "exp", [x]); _}] ); _ - } ] ) - when is_int 1 y && not preserve_stability -> - FunApp (StanLib, "log1m_exp", [x]) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "Minus__" - , [y; {pattern= FunApp (StanLib, "inv_logit", [x]); _}] - ); _ } ] ) - when is_int 1 y && not preserve_stability -> - FunApp (StanLib, "log1m_inv_logit", [x]) - | "log", [{pattern= FunApp (StanLib, "Minus__", [y; x]); _}] - when is_int 1 y && not preserve_stability -> - FunApp (StanLib, "log1m", [x]) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [y; {pattern= FunApp (StanLib, "exp", [x]); _}] ); _ - } ] ) - when is_int 1 y && not preserve_stability -> - FunApp (StanLib, "log1p_exp", [x]) - | "log", [{pattern= FunApp (StanLib, "Plus__", [y; x]); _}] - when is_int 1 y && not preserve_stability -> - FunApp (StanLib, "log1p", [x]) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "fabs" - , [{pattern= FunApp (StanLib, "determinant", [x]); _}] - ); _ } ] ) -> - FunApp (StanLib, "log_determinant", [x]) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "Minus__" - , [ {pattern= FunApp (StanLib, "exp", [x]); _} - ; {pattern= FunApp (StanLib, "exp", [y]); _} ] ); _ - } ] ) -> - FunApp (StanLib, "log_diff_exp", [x; y]) - (* TODO: log_mix?*) - | "log", [{pattern= FunApp (StanLib, "falling_factorial", l); _}] - -> - FunApp (StanLib, "log_falling_factorial", l) - | "log", [{pattern= FunApp (StanLib, "rising_factorial", l); _}] -> - FunApp (StanLib, "log_rising_factorial", l) - | "log", [{pattern= FunApp (StanLib, "inv_logit", l); _}] -> - FunApp (StanLib, "log_inv_logit", l) - | "log", [{pattern= FunApp (StanLib, "softmax", l); _}] -> - FunApp (StanLib, "log_softmax", l) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "sum" - , [{pattern= FunApp (StanLib, "exp", l); _}] ); _ } ] ) - -> - FunApp (StanLib, "log_sum_exp", l) - | ( "log" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ {pattern= FunApp (StanLib, "exp", [x]); _} - ; {pattern= FunApp (StanLib, "exp", [y]); _} ] ); _ - } ] ) -> - FunApp (StanLib, "log_sum_exp", [x; y]) - | ( "multi_normal_lpdf" - , [y; mu; {pattern= FunApp (StanLib, "inverse", [tau]); _}] ) -> - FunApp (StanLib, "multi_normal_prec_lpdf", [y; mu; tau]) - | ( "neg_binomial_2_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } ] ); _ } ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; alpha; beta; sigma] ) - | ( "neg_binomial_2_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } - ; alpha ] ); _ } ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; alpha; beta; sigma] ) - | ( "neg_binomial_2_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [{pattern= FunApp (StanLib, "Times__", [x; beta]); _}] - ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta; sigma] ) - | ( "neg_binomial_2_log_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; alpha; beta; sigma] ) - | ( "neg_binomial_2_log_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ; alpha ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; alpha; beta; sigma] ) - | ( "neg_binomial_2_log_lpmf" - , [y; {pattern= FunApp (StanLib, "Times__", [x; beta]); _}; sigma] - ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "neg_binomial_2_log_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta; sigma] ) - | ( "neg_binomial_2_lpmf" - , [y; {pattern= FunApp (StanLib, "exp", [eta]); _}; phi] ) -> - FunApp (StanLib, "neg_binomial_2_log_lpmf", [y; eta; phi]) - | ( "neg_binomial_2_rng" - , [{pattern= FunApp (StanLib, "exp", [eta]); _}; phi] ) -> - FunApp (StanLib, "neg_binomial_2_log_rng", [eta; phi]) - | ( "normal_lpdf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "normal_id_glm_lpdf", [y; x; alpha; beta; sigma]) - | ( "normal_lpdf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ; alpha ] ); _ } - ; sigma ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - (StanLib, "normal_id_glm_lpdf", [y; x; alpha; beta; sigma]) - | ( "normal_lpdf" - , [y; {pattern= FunApp (StanLib, "Times__", [x; beta]); _}; sigma] - ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "normal_id_glm_lpdf" - , [y; x; Expr.Helpers.zero; beta; sigma] ) - | ( "poisson_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } ] ); _ } ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp (StanLib, "poisson_log_glm_lpmf", [y; x; alpha; beta]) - | ( "poisson_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [ { pattern= - FunApp - ( StanLib - , "Plus__" - , [ { pattern= - FunApp (StanLib, "Times__", [x; beta]); _ - } - ; alpha ] ); _ } ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp (StanLib, "poisson_log_glm_lpmf", [y; x; alpha; beta]) - | ( "poisson_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "exp" - , [{pattern= FunApp (StanLib, "Times__", [x; beta]); _}] - ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "poisson_log_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta] ) - | ( "poisson_log_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ alpha - ; {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp (StanLib, "poisson_log_glm_lpmf", [y; x; alpha; beta]) - | ( "poisson_log_lpmf" - , [ y - ; { pattern= - FunApp - ( StanLib - , "Plus__" - , [ {pattern= FunApp (StanLib, "Times__", [x; beta]); _} - ; alpha ] ); _ } ] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp (StanLib, "poisson_log_glm_lpmf", [y; x; alpha; beta]) - | ( "poisson_log_lpmf" - , [y; {pattern= FunApp (StanLib, "Times__", [x; beta]); _}] ) - when Expr.Typed.type_of x = UMatrix -> - FunApp - ( StanLib - , "poisson_log_glm_lpmf" - , [y; x; Expr.Helpers.zero; beta] ) - | "poisson_lpmf", [y; {pattern= FunApp (StanLib, "exp", [eta]); _}] - -> - FunApp (StanLib, "poisson_log_lpmf", [y; eta]) - | "poisson_rng", [{pattern= FunApp (StanLib, "exp", [eta]); _}] -> - FunApp (StanLib, "poisson_log_rng", [eta]) - | "pow", [y; x] when is_int 2 y -> FunApp (StanLib, "exp2", [x]) - | "rows_dot_product", [x; y] when Expr.Typed.equal x y -> - FunApp (StanLib, "rows_dot_self", [x]) - | "pow", [x; {pattern= Lit (Int, "2"); _}] -> - FunApp (StanLib, "square", [x]) - | "pow", [x; {pattern= Lit (Real, "0.5"); _}] -> - FunApp (StanLib, "sqrt", [x]) - | "pow", [x; {pattern= FunApp (StanLib, "Divide__", [y; z]); _}] - when is_int 1 y && is_int 2 z -> - FunApp (StanLib, "sqrt", [x]) - (* This is wrong; if both are type UInt the exponent is rounds down to zero. *) - | "square", [{pattern= FunApp (StanLib, "sd", [x]); _}] -> - FunApp (StanLib, "variance", [x]) - | "sqrt", [x] when is_int 2 x -> FunApp (StanLib, "sqrt2", []) - | ( "sum" - , [ { pattern= - FunApp - ( StanLib - , "square" - , [{pattern= FunApp (StanLib, "Minus__", [x; y]); _}] ); _ - } ] ) -> - FunApp (StanLib, "squared_distance", [x; y]) - | "sum", [{pattern= FunApp (StanLib, "diagonal", l); _}] -> - FunApp (StanLib, "trace", l) - | ( "trace" - , [ { pattern= - FunApp - ( StanLib - , "Times__" - , [ { pattern= - FunApp - ( StanLib - , "Times__" - , [ { pattern= - FunApp - ( StanLib - , "Times__" - , [ d - ; { pattern= - FunApp - ( StanLib - , "transpose" - , [b] ); _ } ] ); _ } - ; a ] ); _ } - ; c ] ); _ } ] ) - when Expr.Typed.equal b c -> - FunApp (StanLib, "trace_gen_quad_form", [d; a; b]) - | "trace", [{pattern= FunApp (StanLib, "quad_form", [a; b]); _}] -> - FunApp (StanLib, "trace_quad_form", [a; b]) - | "Minus__", [x; {pattern= FunApp (StanLib, "erf", l); _}] - when is_int 1 x -> - FunApp (StanLib, "erfc", l) - | "Minus__", [x; {pattern= FunApp (StanLib, "erfc", l); _}] - when is_int 1 x -> - FunApp (StanLib, "erf", l) - | "Minus__", [{pattern= FunApp (StanLib, "exp", l'); _}; x] - when is_int 1 x && not preserve_stability -> - FunApp (StanLib, "expm1", l') - | "Plus__", [{pattern= FunApp (StanLib, "Times__", [x; y]); _}; z] - |"Plus__", [z; {pattern= FunApp (StanLib, "Times__", [x; y]); _}] - when not preserve_stability -> - FunApp (StanLib, "fma", [x; y; z]) - | "Minus__", [x; {pattern= FunApp (StanLib, "gamma_p", l); _}] - when is_int 1 x -> - FunApp (StanLib, "gamma_q", l) - | "Minus__", [x; {pattern= FunApp (StanLib, "gamma_q", l); _}] - when is_int 1 x -> - FunApp (StanLib, "gamma_p", l) - | ( "Times__" - , [ { pattern= - FunApp - ( StanLib - , "matrix_exp" - , [{pattern= FunApp (StanLib, "Times__", [t; a]); _}] ); _ - } - ; b ] ) - when Expr.Typed.type_of t = UInt || Expr.Typed.type_of t = UReal - -> - FunApp (StanLib, "scale_matrix_exp_multiply", [t; a; b]) - | ( "Times__" - , [ { pattern= - FunApp - ( StanLib - , "matrix_exp" - , [{pattern= FunApp (StanLib, "Times__", [a; t]); _}] ); _ - } - ; b ] ) - when Expr.Typed.type_of t = UInt || Expr.Typed.type_of t = UReal - -> - FunApp (StanLib, "scale_matrix_exp_multiply", [t; a; b]) - | "Times__", [{pattern= FunApp (StanLib, "matrix_exp", [a]); _}; b] - -> - FunApp (StanLib, "matrix_exp_multiply", [a; b]) - | "Times__", [x; {pattern= FunApp (StanLib, "log", [y]); _}] - |"Times__", [{pattern= FunApp (StanLib, "log", [y]); _}; x] - when not preserve_stability -> - FunApp (StanLib, "lmultiply", [x; y]) - | ( "Times__" - , [ {pattern= FunApp (StanLib, "diag_matrix", [v]); _} - ; {pattern= FunApp (StanLib, "diag_post_multiply", [a; w]); _} - ] ) - when Expr.Typed.equal v w -> - FunApp (StanLib, "quad_form_diag", [a; v]) - | ( "Times__" - , [ {pattern= FunApp (StanLib, "diag_pre_multiply", [v; a]); _} - ; {pattern= FunApp (StanLib, "diag_matrix", [w]); _} ] ) - when Expr.Typed.equal v w -> - FunApp (StanLib, "quad_form_diag", [a; v]) - | ( "Times__" - , [ {pattern= FunApp (StanLib, "transpose", [b]); _} - ; {pattern= FunApp (StanLib, "Times__", [a; c]); _} ] ) - when Expr.Typed.equal b c -> - FunApp (StanLib, "quad_form", [a; b]) - | ( "Times__" - , [ { pattern= - FunApp - ( StanLib - , "Times__" - , [{pattern= FunApp (StanLib, "transpose", [b]); _}; a] - ); _ } - ; c ] ) - when Expr.Typed.equal b c -> - FunApp (StanLib, "quad_form", [a; b]) - | ( "Times__" - , [e1'; {pattern= FunApp (StanLib, "diag_matrix", [v]); _}] ) -> - FunApp (StanLib, "diag_post_multiply", [e1'; v]) - | ( "Times__" - , [{pattern= FunApp (StanLib, "diag_matrix", [v]); _}; e2'] ) -> - FunApp (StanLib, "diag_pre_multiply", [v; e2']) - (* Constant folding for operators *) - | op, [{pattern= Lit (Int, i); _}] -> ( - match op with - | "PPlus__" | "PMinus__" | "PNot__" -> - apply_prefix_operator_int op (Int.of_string i) - | _ -> FunApp (t, op, l) ) - | op, [{pattern= Lit (Real, r); _}] -> ( - match op with - | "PPlus__" | "PMinus__" -> - apply_prefix_operator_real op (Float.of_string r) - | _ -> FunApp (t, op, l) ) - | op, [{pattern= Lit (Int, i1); _}; {pattern= Lit (Int, i2); _}] -> ( - match op with - | "Plus__" | "Minus__" | "Times__" | "Divide__" | "IntDivide__" - |"Modulo__" | "Or__" | "And__" | "Equals__" | "NEquals__" - |"Less__" | "Leq__" | "Greater__" | "Geq__" -> - apply_operator_int op (Int.of_string i1) (Int.of_string i2) - | _ -> FunApp (t, op, l) ) - | op, [{pattern= Lit (Real, i1); _}; {pattern= Lit (Real, i2); _}] - |op, [{pattern= Lit (Int, i1); _}; {pattern= Lit (Real, i2); _}] - |op, [{pattern= Lit (Real, i1); _}; {pattern= Lit (Int, i2); _}] - -> ( - match op with - | "Plus__" | "Minus__" | "Times__" | "Divide__" -> - apply_arithmetic_operator_real op (Float.of_string i1) - (Float.of_string i2) - | "Or__" | "And__" | "Equals__" | "NEquals__" | "Less__" - |"Leq__" | "Greater__" | "Geq__" -> - apply_logical_operator_real op (Float.of_string i1) - (Float.of_string i2) - | _ -> FunApp (t, op, l) ) - | _ -> FunApp (t, f, l) ) - | TernaryIf (e1, e2, e3) -> ( - match (eval_expr e1, eval_expr e2, eval_expr e3) with + Operator.of_string_opt name + |> Option.value_map + ~f:(fun op -> + Stan_math_signatures.operator_stan_math_return_type op argument_types) + ~default:(Stan_math_signatures.stan_math_returntype name argument_types) + in + let try_partially_evaluate_to e = + Expr.Fixed.Pattern.( + match e with + | FunApp (StanLib, f', l') -> + (match get_fun_or_op_rt_opt f' l' with + | Some _ -> FunApp (StanLib, f', l') + | None -> FunApp (StanLib, f, l)) + | e -> e) + in + try_partially_evaluate_to + (match f, l with + (* TODO: deal with tilde statements and unnormalized distributions properly here *) + | ( "bernoulli_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "inv_logit" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ] ) + ; _ + } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; alpha; beta ]) + | ( "bernoulli_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "inv_logit" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ; alpha + ] ) + ; _ + } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; alpha; beta ]) + | ( "bernoulli_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "inv_logit" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; Expr.Helpers.zero; beta ]) + | ( "bernoulli_logit_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; alpha; beta ]) + | ( "bernoulli_logit_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ; alpha + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; alpha; beta ]) + | ( "bernoulli_logit_lpmf" + , [ y; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "bernoulli_logit_glm_lpmf", [ y; x; Expr.Helpers.zero; beta ]) + | ( "bernoulli_lpmf" + , [ y; { pattern = FunApp (StanLib, "inv_logit", [ alpha ]); _ } ] ) -> + FunApp (StanLib, "bernoulli_logit_lpmf", [ y; alpha ]) + | "bernoulli_rng", [ { pattern = FunApp (StanLib, "inv_logit", [ alpha ]); _ } ] + -> FunApp (StanLib, "bernoulli_logit_rng", [ alpha ]) + | ( "binomial_lpmf" + , [ y; n; { pattern = FunApp (StanLib, "inv_logit", [ alpha ]); _ } ] ) -> + FunApp (StanLib, "binomial_logit_lpmf", [ y; n; alpha ]) + | ( "categorical_lpmf" + , [ y; { pattern = FunApp (StanLib, "inv_logit", [ alpha ]); _ } ] ) -> + FunApp (StanLib, "categorical_logit_lpmf", [ y; alpha ]) + | ( "categorical_rng" + , [ { pattern = FunApp (StanLib, "inv_logit", [ alpha ]); _ } ] ) -> + FunApp (StanLib, "categorical_logit_rng", [ alpha ]) + | "columns_dot_product", [ x; y ] when Expr.Typed.equal x y -> + FunApp (StanLib, "columns_dot_self", [ x ]) + | "dot_product", [ x; y ] when Expr.Typed.equal x y -> + FunApp (StanLib, "dot_self", [ x ]) + | "inv", [ { pattern = FunApp (StanLib, "sqrt", l); _ } ] -> + FunApp (StanLib, "inv_sqrt", l) + | "inv", [ { pattern = FunApp (StanLib, "square", [ x ]); _ } ] -> + FunApp (StanLib, "inv_square", [ x ]) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "Minus__" + , [ y; { pattern = FunApp (StanLib, "exp", [ x ]); _ } ] ) + ; _ + } + ] ) + when is_int 1 y && not preserve_stability -> + FunApp (StanLib, "log1m_exp", [ x ]) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "Minus__" + , [ y; { pattern = FunApp (StanLib, "inv_logit", [ x ]); _ } ] ) + ; _ + } + ] ) + when is_int 1 y && not preserve_stability -> + FunApp (StanLib, "log1m_inv_logit", [ x ]) + | "log", [ { pattern = FunApp (StanLib, "Minus__", [ y; x ]); _ } ] + when is_int 1 y && not preserve_stability -> FunApp (StanLib, "log1m", [ x ]) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ y; { pattern = FunApp (StanLib, "exp", [ x ]); _ } ] ) + ; _ + } + ] ) + when is_int 1 y && not preserve_stability -> + FunApp (StanLib, "log1p_exp", [ x ]) + | "log", [ { pattern = FunApp (StanLib, "Plus__", [ y; x ]); _ } ] + when is_int 1 y && not preserve_stability -> FunApp (StanLib, "log1p", [ x ]) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "fabs" + , [ { pattern = FunApp (StanLib, "determinant", [ x ]); _ } ] ) + ; _ + } + ] ) -> FunApp (StanLib, "log_determinant", [ x ]) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "Minus__" + , [ { pattern = FunApp (StanLib, "exp", [ x ]); _ } + ; { pattern = FunApp (StanLib, "exp", [ y ]); _ } + ] ) + ; _ + } + ] ) -> FunApp (StanLib, "log_diff_exp", [ x; y ]) + (* TODO: log_mix?*) + | "log", [ { pattern = FunApp (StanLib, "falling_factorial", l); _ } ] -> + FunApp (StanLib, "log_falling_factorial", l) + | "log", [ { pattern = FunApp (StanLib, "rising_factorial", l); _ } ] -> + FunApp (StanLib, "log_rising_factorial", l) + | "log", [ { pattern = FunApp (StanLib, "inv_logit", l); _ } ] -> + FunApp (StanLib, "log_inv_logit", l) + | "log", [ { pattern = FunApp (StanLib, "softmax", l); _ } ] -> + FunApp (StanLib, "log_softmax", l) + | ( "log" + , [ { pattern = + FunApp + (StanLib, "sum", [ { pattern = FunApp (StanLib, "exp", l); _ } ]) + ; _ + } + ] ) -> FunApp (StanLib, "log_sum_exp", l) + | ( "log" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "exp", [ x ]); _ } + ; { pattern = FunApp (StanLib, "exp", [ y ]); _ } + ] ) + ; _ + } + ] ) -> FunApp (StanLib, "log_sum_exp", [ x; y ]) + | ( "multi_normal_lpdf" + , [ y; mu; { pattern = FunApp (StanLib, "inverse", [ tau ]); _ } ] ) -> + FunApp (StanLib, "multi_normal_prec_lpdf", [ y; mu; tau ]) + | ( "neg_binomial_2_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ] ) + ; _ + } + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "neg_binomial_2_log_glm_lpmf", [ y; x; alpha; beta; sigma ]) + | ( "neg_binomial_2_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ; alpha + ] ) + ; _ + } + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "neg_binomial_2_log_glm_lpmf", [ y; x; alpha; beta; sigma ]) + | ( "neg_binomial_2_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp + ( StanLib + , "neg_binomial_2_log_glm_lpmf" + , [ y; x; Expr.Helpers.zero; beta; sigma ] ) + | ( "neg_binomial_2_log_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "neg_binomial_2_log_glm_lpmf", [ y; x; alpha; beta; sigma ]) + | ( "neg_binomial_2_log_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ; alpha + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "neg_binomial_2_log_glm_lpmf", [ y; x; alpha; beta; sigma ]) + | ( "neg_binomial_2_log_lpmf" + , [ y; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ }; sigma ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp + ( StanLib + , "neg_binomial_2_log_glm_lpmf" + , [ y; x; Expr.Helpers.zero; beta; sigma ] ) + | ( "neg_binomial_2_lpmf" + , [ y; { pattern = FunApp (StanLib, "exp", [ eta ]); _ }; phi ] ) -> + FunApp (StanLib, "neg_binomial_2_log_lpmf", [ y; eta; phi ]) + | ( "neg_binomial_2_rng" + , [ { pattern = FunApp (StanLib, "exp", [ eta ]); _ }; phi ] ) -> + FunApp (StanLib, "neg_binomial_2_log_rng", [ eta; phi ]) + | ( "normal_lpdf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "normal_id_glm_lpdf", [ y; x; alpha; beta; sigma ]) + | ( "normal_lpdf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ; alpha + ] ) + ; _ + } + ; sigma + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "normal_id_glm_lpdf", [ y; x; alpha; beta; sigma ]) + | ( "normal_lpdf" + , [ y; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ }; sigma ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp + (StanLib, "normal_id_glm_lpdf", [ y; x; Expr.Helpers.zero; beta; sigma ]) + | ( "poisson_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ] ) + ; _ + } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; alpha; beta ]) + | ( "poisson_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]) + ; _ + } + ; alpha + ] ) + ; _ + } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; alpha; beta ]) + | ( "poisson_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "exp" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; Expr.Helpers.zero; beta ]) + | ( "poisson_log_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ alpha + ; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; alpha; beta ]) + | ( "poisson_log_lpmf" + , [ y + ; { pattern = + FunApp + ( StanLib + , "Plus__" + , [ { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } + ; alpha + ] ) + ; _ + } + ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; alpha; beta ]) + | ( "poisson_log_lpmf" + , [ y; { pattern = FunApp (StanLib, "Times__", [ x; beta ]); _ } ] ) + when Expr.Typed.type_of x = UMatrix -> + FunApp (StanLib, "poisson_log_glm_lpmf", [ y; x; Expr.Helpers.zero; beta ]) + | "poisson_lpmf", [ y; { pattern = FunApp (StanLib, "exp", [ eta ]); _ } ] -> + FunApp (StanLib, "poisson_log_lpmf", [ y; eta ]) + | "poisson_rng", [ { pattern = FunApp (StanLib, "exp", [ eta ]); _ } ] -> + FunApp (StanLib, "poisson_log_rng", [ eta ]) + | "pow", [ y; x ] when is_int 2 y -> FunApp (StanLib, "exp2", [ x ]) + | "rows_dot_product", [ x; y ] when Expr.Typed.equal x y -> + FunApp (StanLib, "rows_dot_self", [ x ]) + | "pow", [ x; { pattern = Lit (Int, "2"); _ } ] -> + FunApp (StanLib, "square", [ x ]) + | "pow", [ x; { pattern = Lit (Real, "0.5"); _ } ] -> + FunApp (StanLib, "sqrt", [ x ]) + | "pow", [ x; { pattern = FunApp (StanLib, "Divide__", [ y; z ]); _ } ] + when is_int 1 y && is_int 2 z -> + FunApp (StanLib, "sqrt", [ x ]) + (* This is wrong; if both are type UInt the exponent is rounds down to zero. *) + | "square", [ { pattern = FunApp (StanLib, "sd", [ x ]); _ } ] -> + FunApp (StanLib, "variance", [ x ]) + | "sqrt", [ x ] when is_int 2 x -> FunApp (StanLib, "sqrt2", []) + | ( "sum" + , [ { pattern = + FunApp + ( StanLib + , "square" + , [ { pattern = FunApp (StanLib, "Minus__", [ x; y ]); _ } ] ) + ; _ + } + ] ) -> FunApp (StanLib, "squared_distance", [ x; y ]) + | "sum", [ { pattern = FunApp (StanLib, "diagonal", l); _ } ] -> + FunApp (StanLib, "trace", l) + | ( "trace" + , [ { pattern = + FunApp + ( StanLib + , "Times__" + , [ { pattern = + FunApp + ( StanLib + , "Times__" + , [ { pattern = + FunApp + ( StanLib + , "Times__" + , [ d + ; { pattern = + FunApp (StanLib, "transpose", [ b ]) + ; _ + } + ] ) + ; _ + } + ; a + ] ) + ; _ + } + ; c + ] ) + ; _ + } + ] ) + when Expr.Typed.equal b c -> + FunApp (StanLib, "trace_gen_quad_form", [ d; a; b ]) + | "trace", [ { pattern = FunApp (StanLib, "quad_form", [ a; b ]); _ } ] -> + FunApp (StanLib, "trace_quad_form", [ a; b ]) + | "Minus__", [ x; { pattern = FunApp (StanLib, "erf", l); _ } ] when is_int 1 x + -> FunApp (StanLib, "erfc", l) + | "Minus__", [ x; { pattern = FunApp (StanLib, "erfc", l); _ } ] when is_int 1 x + -> FunApp (StanLib, "erf", l) + | "Minus__", [ { pattern = FunApp (StanLib, "exp", l'); _ }; x ] + when is_int 1 x && not preserve_stability -> FunApp (StanLib, "expm1", l') + | "Plus__", [ { pattern = FunApp (StanLib, "Times__", [ x; y ]); _ }; z ] + | "Plus__", [ z; { pattern = FunApp (StanLib, "Times__", [ x; y ]); _ } ] + when not preserve_stability -> FunApp (StanLib, "fma", [ x; y; z ]) + | "Minus__", [ x; { pattern = FunApp (StanLib, "gamma_p", l); _ } ] + when is_int 1 x -> FunApp (StanLib, "gamma_q", l) + | "Minus__", [ x; { pattern = FunApp (StanLib, "gamma_q", l); _ } ] + when is_int 1 x -> FunApp (StanLib, "gamma_p", l) + | ( "Times__" + , [ { pattern = + FunApp + ( StanLib + , "matrix_exp" + , [ { pattern = FunApp (StanLib, "Times__", [ t; a ]); _ } ] ) + ; _ + } + ; b + ] ) + when Expr.Typed.type_of t = UInt || Expr.Typed.type_of t = UReal -> + FunApp (StanLib, "scale_matrix_exp_multiply", [ t; a; b ]) + | ( "Times__" + , [ { pattern = + FunApp + ( StanLib + , "matrix_exp" + , [ { pattern = FunApp (StanLib, "Times__", [ a; t ]); _ } ] ) + ; _ + } + ; b + ] ) + when Expr.Typed.type_of t = UInt || Expr.Typed.type_of t = UReal -> + FunApp (StanLib, "scale_matrix_exp_multiply", [ t; a; b ]) + | "Times__", [ { pattern = FunApp (StanLib, "matrix_exp", [ a ]); _ }; b ] -> + FunApp (StanLib, "matrix_exp_multiply", [ a; b ]) + | "Times__", [ x; { pattern = FunApp (StanLib, "log", [ y ]); _ } ] + | "Times__", [ { pattern = FunApp (StanLib, "log", [ y ]); _ }; x ] + when not preserve_stability -> FunApp (StanLib, "lmultiply", [ x; y ]) + | ( "Times__" + , [ { pattern = FunApp (StanLib, "diag_matrix", [ v ]); _ } + ; { pattern = FunApp (StanLib, "diag_post_multiply", [ a; w ]); _ } + ] ) + when Expr.Typed.equal v w -> FunApp (StanLib, "quad_form_diag", [ a; v ]) + | ( "Times__" + , [ { pattern = FunApp (StanLib, "diag_pre_multiply", [ v; a ]); _ } + ; { pattern = FunApp (StanLib, "diag_matrix", [ w ]); _ } + ] ) + when Expr.Typed.equal v w -> FunApp (StanLib, "quad_form_diag", [ a; v ]) + | ( "Times__" + , [ { pattern = FunApp (StanLib, "transpose", [ b ]); _ } + ; { pattern = FunApp (StanLib, "Times__", [ a; c ]); _ } + ] ) + when Expr.Typed.equal b c -> FunApp (StanLib, "quad_form", [ a; b ]) + | ( "Times__" + , [ { pattern = + FunApp + ( StanLib + , "Times__" + , [ { pattern = FunApp (StanLib, "transpose", [ b ]); _ }; a ] ) + ; _ + } + ; c + ] ) + when Expr.Typed.equal b c -> FunApp (StanLib, "quad_form", [ a; b ]) + | "Times__", [ e1'; { pattern = FunApp (StanLib, "diag_matrix", [ v ]); _ } ] -> + FunApp (StanLib, "diag_post_multiply", [ e1'; v ]) + | "Times__", [ { pattern = FunApp (StanLib, "diag_matrix", [ v ]); _ }; e2' ] -> + FunApp (StanLib, "diag_pre_multiply", [ v; e2' ]) + (* Constant folding for operators *) + | op, [ { pattern = Lit (Int, i); _ } ] -> + (match op with + | "PPlus__" | "PMinus__" | "PNot__" -> + apply_prefix_operator_int op (Int.of_string i) + | _ -> FunApp (t, op, l)) + | op, [ { pattern = Lit (Real, r); _ } ] -> + (match op with + | "PPlus__" | "PMinus__" -> apply_prefix_operator_real op (Float.of_string r) + | _ -> FunApp (t, op, l)) + | op, [ { pattern = Lit (Int, i1); _ }; { pattern = Lit (Int, i2); _ } ] -> + (match op with + | "Plus__" + | "Minus__" + | "Times__" + | "Divide__" + | "IntDivide__" + | "Modulo__" + | "Or__" + | "And__" + | "Equals__" + | "NEquals__" + | "Less__" + | "Leq__" + | "Greater__" + | "Geq__" -> apply_operator_int op (Int.of_string i1) (Int.of_string i2) + | _ -> FunApp (t, op, l)) + | op, [ { pattern = Lit (Real, i1); _ }; { pattern = Lit (Real, i2); _ } ] + | op, [ { pattern = Lit (Int, i1); _ }; { pattern = Lit (Real, i2); _ } ] + | op, [ { pattern = Lit (Real, i1); _ }; { pattern = Lit (Int, i2); _ } ] -> + (match op with + | "Plus__" | "Minus__" | "Times__" | "Divide__" -> + apply_arithmetic_operator_real op (Float.of_string i1) (Float.of_string i2) + | "Or__" + | "And__" + | "Equals__" + | "NEquals__" + | "Less__" + | "Leq__" + | "Greater__" + | "Geq__" -> + apply_logical_operator_real op (Float.of_string i1) (Float.of_string i2) + | _ -> FunApp (t, op, l)) + | _ -> FunApp (t, f, l)) + | TernaryIf (e1, e2, e3) -> + (match eval_expr e1, eval_expr e2, eval_expr e3 with | x, _, e3' when is_int 0 x -> e3'.pattern - | {pattern= Lit (Int, _); _}, e2', _ -> e2'.pattern - | e1', e2', e3' -> TernaryIf (e1', e2', e3') ) - | EAnd (e1, e2) -> ( - match (eval_expr e1, eval_expr e2) with - | {pattern= Lit (Int, s1); _}, {pattern= Lit (Int, s2); _} -> - let i1, i2 = (Int.of_string s1, Int.of_string s2) in - Lit (Int, Int.to_string (Bool.to_int (i1 <> 0 && i2 <> 0))) - | {pattern= Lit (_, s1); _}, {pattern= Lit (_, s2); _} -> - let r1, r2 = (Float.of_string s1, Float.of_string s2) in - Lit (Int, Int.to_string (Bool.to_int (r1 <> 0. && r2 <> 0.))) - | e1', e2' -> EAnd (e1', e2') ) - | EOr (e1, e2) -> ( - match (eval_expr e1, eval_expr e2) with - | {pattern= Lit (Int, s1); _}, {pattern= Lit (Int, s2); _} -> - let i1, i2 = (Int.of_string s1, Int.of_string s2) in - Lit (Int, Int.to_string (Bool.to_int (i1 <> 0 || i2 <> 0))) - | {pattern= Lit (_, s1); _}, {pattern= Lit (_, s2); _} -> - let r1, r2 = (Float.of_string s1, Float.of_string s2) in - Lit (Int, Int.to_string (Bool.to_int (r1 <> 0. || r2 <> 0.))) - | e1', e2' -> EOr (e1', e2') ) + | { pattern = Lit (Int, _); _ }, e2', _ -> e2'.pattern + | e1', e2', e3' -> TernaryIf (e1', e2', e3')) + | EAnd (e1, e2) -> + (match eval_expr e1, eval_expr e2 with + | { pattern = Lit (Int, s1); _ }, { pattern = Lit (Int, s2); _ } -> + let i1, i2 = Int.of_string s1, Int.of_string s2 in + Lit (Int, Int.to_string (Bool.to_int (i1 <> 0 && i2 <> 0))) + | { pattern = Lit (_, s1); _ }, { pattern = Lit (_, s2); _ } -> + let r1, r2 = Float.of_string s1, Float.of_string s2 in + Lit (Int, Int.to_string (Bool.to_int (r1 <> 0. && r2 <> 0.))) + | e1', e2' -> EAnd (e1', e2')) + | EOr (e1, e2) -> + (match eval_expr e1, eval_expr e2 with + | { pattern = Lit (Int, s1); _ }, { pattern = Lit (Int, s2); _ } -> + let i1, i2 = Int.of_string s1, Int.of_string s2 in + Lit (Int, Int.to_string (Bool.to_int (i1 <> 0 || i2 <> 0))) + | { pattern = Lit (_, s1); _ }, { pattern = Lit (_, s2); _ } -> + let r1, r2 = Float.of_string s1, Float.of_string s2 in + Lit (Int, Int.to_string (Bool.to_int (r1 <> 0. || r2 <> 0.))) + | e1', e2' -> EOr (e1', e2')) | Indexed (e, l) -> - (* TODO: do something clever with array and matrix expressions here? - Note that we could also constant fold array sizes if we keep those around on declarations. *) - Indexed (eval_expr e, List.map ~f:(Index.map eval_expr) l) ) } + (* TODO: do something clever with array and matrix expressions here? + Note that we could also constant fold array sizes if we keep those around on declarations. *) + Indexed (eval_expr e, List.map ~f:(Index.map eval_expr) l)) + } +;; let rec simplify_index_expr pattern = Expr.Fixed.( match pattern with | Pattern.Indexed - ( { pattern= + ( { pattern = Indexed (obj, inner_indices) (* , Single ({emeta= {type_= UArray UInt; _} as emeta; _} as multi) * :: inner_tl ) *) - ; meta } - , ( Single ({meta= Expr.Typed.Meta.({type_= UInt; _}); _} as single_e) - as single ) + ; meta + } + , (Single ({ meta = Expr.Typed.Meta.{ type_ = UInt; _ }; _ } as single_e) as + single) :: outer_tl ) - when List.exists ~f:is_multi_index inner_indices -> ( - match List.split_while ~f:(Fn.non is_multi_index) inner_indices with + when List.exists ~f:is_multi_index inner_indices -> + (match List.split_while ~f:(Fn.non is_multi_index) inner_indices with | inner_singles, MultiIndex first_multi :: inner_tl -> - (* foo [arr1, ..., arrN] [i1, ..., iN] -> - foo [arr1[i1]] [arr[i2]] ... [arrN[iN]] *) - simplify_index_expr - (Indexed - ( { pattern= - Indexed - ( obj - , inner_singles - @ [ Index.Single - { pattern= Indexed (first_multi, [single]) - ; meta= {meta with type_= UInt} } ] - @ inner_tl ) - ; meta } - , outer_tl )) + (* foo [arr1, ..., arrN] [i1, ..., iN] -> + foo [arr1[i1]] [arr[i2]] ... [arrN[iN]] *) + simplify_index_expr + (Indexed + ( { pattern = + Indexed + ( obj + , inner_singles + @ [ Index.Single + { pattern = Indexed (first_multi, [ single ]) + ; meta = { meta with type_ = UInt } + } + ] + @ inner_tl ) + ; meta + } + , outer_tl )) | inner_singles, All :: inner_tl -> - (* v[:x][i] -> v[i] *) - (* v[:][i] -> v[i] *) - (* XXX generate check *) - simplify_index_expr - (Indexed - ( { pattern= Indexed (obj, inner_singles @ [single] @ inner_tl) - ; meta } - , outer_tl )) + (* v[:x][i] -> v[i] *) + (* v[:][i] -> v[i] *) + (* XXX generate check *) + simplify_index_expr + (Indexed + ( { pattern = Indexed (obj, inner_singles @ [ single ] @ inner_tl); meta } + , outer_tl )) | inner_singles, Between (bot, _) :: inner_tl - |inner_singles, Upfrom bot :: inner_tl -> - (* v[x:y][z] -> v[x+z-1] *) - (* XXX generate check *) - simplify_index_expr - (Indexed - ( { pattern= - Indexed - ( obj - , inner_singles - @ [ Index.Single - Expr.Helpers.( - binop (binop bot Plus single_e) Minus - loop_bottom) ] - @ inner_tl ) - ; meta } - , outer_tl )) + | inner_singles, Upfrom bot :: inner_tl -> + (* v[x:y][z] -> v[x+z-1] *) + (* XXX generate check *) + simplify_index_expr + (Indexed + ( { pattern = + Indexed + ( obj + , inner_singles + @ [ Index.Single + Expr.Helpers.( + binop (binop bot Plus single_e) Minus loop_bottom) + ] + @ inner_tl ) + ; meta + } + , outer_tl )) | inner_singles, (([] | Single _ :: _) as multis) -> - raise_s - [%message - "Impossible! There must be a multi-index." - (inner_singles : Expr.Typed.t Index.t list) - (multis : Expr.Typed.t Index.t list)] ) + raise_s + [%message + "Impossible! There must be a multi-index." + (inner_singles : Expr.Typed.t Index.t list) + (multis : Expr.Typed.t Index.t list)]) | e -> e) +;; let remove_trailing_alls_expr = function | Expr.Fixed.Pattern.Indexed (obj, indices) -> - (* a[2][:] -> a[2] *) - let rec remove_trailing_alls indices = - match List.rev indices with - | Index.All :: tl -> remove_trailing_alls (List.rev tl) - | _ -> indices - in - Expr.Fixed.Pattern.Indexed (obj, remove_trailing_alls indices) + (* a[2][:] -> a[2] *) + let rec remove_trailing_alls indices = + match List.rev indices with + | Index.All :: tl -> remove_trailing_alls (List.rev tl) + | _ -> indices + in + Expr.Fixed.Pattern.Indexed (obj, remove_trailing_alls indices) | e -> e +;; let rec simplify_indices_expr expr = Expr.Fixed.( let pattern = - expr.pattern |> remove_trailing_alls_expr |> simplify_index_expr + expr.pattern + |> remove_trailing_alls_expr + |> simplify_index_expr |> Expr.Fixed.Pattern.map simplify_indices_expr in - {expr with pattern}) + { expr with pattern }) +;; let eval_stmt_base = Stmt.Fixed.Pattern.map (Fn.compose eval_expr simplify_indices_expr) Fn.id +;; let eval_stmt = map_rec_stmt_loc eval_stmt_base let eval_prog = Program.map eval_expr eval_stmt diff --git a/src/analysis_and_optimization/Pedantic_analysis.ml b/src/analysis_and_optimization/Pedantic_analysis.ml index 49a8cd94a1..43b27e3679 100644 --- a/src/analysis_and_optimization/Pedantic_analysis.ml +++ b/src/analysis_and_optimization/Pedantic_analysis.ml @@ -13,8 +13,9 @@ open Pedantic_dist_warnings Pattern collection functions ********************) -let list_unused_params (factor_graph : factor_graph) (mir : Program.Typed.t) : - string Set.Poly.t = +let list_unused_params (factor_graph : factor_graph) (mir : Program.Typed.t) + : string Set.Poly.t + = (* Build a factor graph of the program, check for missing parameters *) let params = parameter_names_set ~include_transformed:true mir in let used_params = @@ -23,105 +24,115 @@ let list_unused_params (factor_graph : factor_graph) (mir : Program.Typed.t) : (Set.Poly.of_list (Map.Poly.keys factor_graph.var_map)) in Set.Poly.diff params used_params +;; -let list_hard_constrained (mir : Program.Typed.t) : - (string * [`HardConstraint | `NonsenseConstraint]) Set.Poly.t = +let list_hard_constrained (mir : Program.Typed.t) + : (string * [ `HardConstraint | `NonsenseConstraint ]) Set.Poly.t + = (* Iterate through all parameters' transformations for hard constraints *) let constrained (e : bound_values) = match e with - | {lower= `Lit 0.; upper= `Lit 1.} | {lower= `Lit -1.; upper= `Lit 1.} -> - None - | {lower= `Lit a; upper= `Lit b} when a >= b -> Some `NonsenseConstraint - | {lower= `Lit _; upper= `Lit _} -> Some `HardConstraint + | { lower = `Lit 0.; upper = `Lit 1. } | { lower = `Lit -1.; upper = `Lit 1. } -> None + | { lower = `Lit a; upper = `Lit b } when a >= b -> Some `NonsenseConstraint + | { lower = `Lit _; upper = `Lit _ } -> Some `HardConstraint | _ -> None in Set.Poly.filter_map ~f:(fun (name, trans) -> - Option.map - ~f:(fun c -> (name, c)) - (constrained (trans_bounds_values trans)) ) + Option.map ~f:(fun c -> name, c) (constrained (trans_bounds_values trans))) (parameter_set mir) +;; -let list_multi_twiddles (mir : Program.Typed.t) : - (string * Location_span.t Set.Poly.t) Set.Poly.t = +let list_multi_twiddles (mir : Program.Typed.t) + : (string * Location_span.t Set.Poly.t) Set.Poly.t + = (* Collect statements of the form "target += Dist(param, ...)" *) - let collect_twiddle_stmt (stmt : Stmt.Located.t) : - (string, Location_span.t Set.Poly.t) Map.Poly.t = + let collect_twiddle_stmt (stmt : Stmt.Located.t) + : (string, Location_span.t Set.Poly.t) Map.Poly.t + = match stmt.pattern with | Stmt.Fixed.Pattern.TargetPE - { pattern= - Expr.Fixed.Pattern.FunApp (_, _, {pattern= Var vname; _} :: _); _ - } -> - Map.Poly.singleton vname (Set.Poly.singleton stmt.meta) + { pattern = Expr.Fixed.Pattern.FunApp (_, _, { pattern = Var vname; _ } :: _); _ } + -> Map.Poly.singleton vname (Set.Poly.singleton stmt.meta) | _ -> Map.Poly.empty in let twiddles = fold_stmts ~take_stmt:(fun m s -> merge_set_maps m (collect_twiddle_stmt s)) ~take_expr:(fun m _ -> m) - ~init:Map.Poly.empty mir.log_prob + ~init:Map.Poly.empty + mir.log_prob in (* Filter for parameters assigned more than one distribution *) - let multi_twiddles = - Map.Poly.filter ~f:(fun s -> Set.Poly.length s <> 1) twiddles - in - Map.fold ~init:Set.Poly.empty + let multi_twiddles = Map.Poly.filter ~f:(fun s -> Set.Poly.length s <> 1) twiddles in + Map.fold + ~init:Set.Poly.empty ~f:(fun ~key ~data s -> Set.add s (key, data)) multi_twiddles +;; (* Find all of the targets which are dependencies for a given label *) -let var_deps info_map label ?expr:(expr_opt : Expr.Typed.t option = None) - (targets : string Set.Poly.t) : string Set.Poly.t = +let var_deps + info_map + label + ?expr:(expr_opt : Expr.Typed.t option = None) + (targets : string Set.Poly.t) + : string Set.Poly.t + = (* Labels of dependencies *) let dep_labels = match expr_opt with | None -> node_dependencies info_map label | Some expr -> - let vars = Set.Poly.map ~f:fst (expr_var_set expr) in - node_vars_dependencies info_map vars label + let vars = Set.Poly.map ~f:fst (expr_var_set expr) in + node_vars_dependencies info_map vars label in (* expressions of dependencies *) let dep_exprs = union_map dep_labels ~f:(fun label -> let stmt, _ = Map.Poly.find_exn info_map label in - stmt_rhs_var_set stmt ) + stmt_rhs_var_set stmt) in (* variable dependencies *) let dep_vars = Set.Poly.map ~f:(fun (VVar v, _) -> v) dep_exprs in (* target dependencies *) Set.Poly.inter targets dep_vars +;; let list_target_dependant_cf (info_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) (targets : string Set.Poly.t) : - (Location_span.t * string Set.Poly.t) Set.Poly.t = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + (targets : string Set.Poly.t) + : (Location_span.t * string Set.Poly.t) Set.Poly.t + = (* Find all the control flow nodes *) let cf_labels = Set.Poly.of_list - (Map.Poly.keys - (Map.Poly.filter info_map ~f:(fun (stmt, _) -> is_ctrl_flow stmt))) + (Map.Poly.keys (Map.Poly.filter info_map ~f:(fun (stmt, _) -> is_ctrl_flow stmt))) in Set.Poly.filter_map ~f:(fun label -> let deps = var_deps info_map label targets in - if Set.Poly.is_empty deps then None - else + if Set.Poly.is_empty deps + then None + else ( let _, info = Map.Poly.find_exn info_map label in - Some (info.meta, deps) ) + Some (info.meta, deps))) cf_labels +;; -let list_param_dependant_cf (mir : Program.Typed.t) : - (Location_span.t * string Set.Poly.t) Set.Poly.t = +let list_param_dependant_cf (mir : Program.Typed.t) + : (Location_span.t * string Set.Poly.t) Set.Poly.t + = let params = parameter_names_set mir in (* build dataflow data structure *) let info_map = log_prob_build_dep_info_map mir in list_target_dependant_cf info_map params +;; -let list_arg_dependant_fundef_cf (mir : Program.Typed.t) - (fun_def : 'a Program.fun_def) : - (Location_span.t * int * string) Set.Poly.t = +let list_arg_dependant_fundef_cf (mir : Program.Typed.t) (fun_def : 'a Program.fun_def) + : (Location_span.t * int * string) Set.Poly.t + = let args = List.map ~f:(fun (_, name, _) -> name) fun_def.fdargs in (* build dataflow data structure *) let info_map = build_dep_info_map mir fun_def.fdbody in @@ -131,122 +142,135 @@ let list_arg_dependant_fundef_cf (mir : Program.Typed.t) let ix, _ = Option.value_exn ~message: - "INTERNAL ERROR: Pedantic mode found CF dependent on an \ - arg,but the arg is mismatched. Please report a bug.\n" + "INTERNAL ERROR: Pedantic mode found CF dependent on an arg,but the arg \ + is mismatched. Please report a bug.\n" (List.findi args ~f:(fun _ arg -> arg = name)) in - (loc, ix, name) ) ) + loc, ix, name)) +;; let expr_collect_exprs (expr : Expr.Typed.t) ~f : 'a Set.Poly.t = let collect_expr s (expr : Expr.Typed.t) = - match f expr with Some a -> Set.Poly.add s a | _ -> s + match f expr with + | Some a -> Set.Poly.add s a + | _ -> s in fold_expr ~init:Set.Poly.empty ~take_expr:(fun s e -> collect_expr s e) expr +;; let stmts_collect_exprs - (stmts : (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t List.t) ~f : - 'a Set.Poly.t = + (stmts : (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t List.t) + ~f + : 'a Set.Poly.t + = let collect_expr s (expr : Expr.Typed.t) = - match f expr with Some a -> Set.Poly.add s a | _ -> s + match f expr with + | Some a -> Set.Poly.add s a + | _ -> s in - fold_stmts ~init:Set.Poly.empty + fold_stmts + ~init:Set.Poly.empty ~take_stmt:(fun s _ -> s) ~take_expr:(fun s e -> collect_expr s e) stmts +;; -let list_param_dependant_fundef_cf (mir : Program.Typed.t) +let list_param_dependant_fundef_cf + (mir : Program.Typed.t) (info_map : - ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info ) - Map.Poly.t) (fun_def : 'a Program.fun_def) : - (Location_span.t * string Set.Poly.t * string * Location_span.t) Set.Poly.t - = + (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * node_dep_info) Map.Poly.t) + (fun_def : 'a Program.fun_def) + : (Location_span.t * string Set.Poly.t * string * Location_span.t) Set.Poly.t + = let dep_args = list_arg_dependant_fundef_cf mir fun_def in let fun_calls : (Expr.Typed.t * label) Set.Poly.t = Set.Poly.union_list - (List.map ~f:snd + (List.map + ~f:snd (Map.Poly.to_alist - (Map.Poly.filter_mapi info_map - ~f:(fun ~key:label ~data:(stmt, _) -> + (Map.Poly.filter_mapi info_map ~f:(fun ~key:label ~data:(stmt, _) -> let funapps = union_map (stmt_rhs stmt) ~f:(fun rhs_expr -> expr_collect_exprs rhs_expr ~f:(fun rhs_subexpr -> match rhs_subexpr.pattern with | Expr.Fixed.Pattern.FunApp (UserDefined, fname, _) - when fname = fun_def.fdname -> - Some (rhs_subexpr, label) - | _ -> None ) ) + when fname = fun_def.fdname -> Some (rhs_subexpr, label) + | _ -> None)) in - if Set.Poly.is_empty funapps then None else Some funapps )))) + if Set.Poly.is_empty funapps then None else Some funapps)))) in let arg_exprs (fcall_expr : Expr.Typed.t) = match fcall_expr with - | {pattern= Expr.Fixed.Pattern.FunApp (UserDefined, fname, arg_exprs); _} + | { pattern = Expr.Fixed.Pattern.FunApp (UserDefined, fname, arg_exprs); _ } when fname = fun_def.fdname -> - Set.Poly.map dep_args ~f:(fun (loc, ix, arg_name) -> - (loc, List.nth_exn arg_exprs ix, arg_name) ) + Set.Poly.map dep_args ~f:(fun (loc, ix, arg_name) -> + loc, List.nth_exn arg_exprs ix, arg_name) | _ -> - raise - (Failure - "In finding searching for parameter dependent functionarguments, \ - mismatched function. Please report a bug.\n") + raise + (Failure + "In finding searching for parameter dependent functionarguments, mismatched \ + function. Please report a bug.\n") in let arg_param_deps label arg_expr = var_deps info_map ~expr:(Some arg_expr) label (parameter_names_set mir) in union_map fun_calls ~f:(fun (fcall_expr, label) -> - Set.Poly.filter_map (arg_exprs fcall_expr) - ~f:(fun (cf_loc, arg_expr, arg_name) -> + Set.Poly.filter_map (arg_exprs fcall_expr) ~f:(fun (cf_loc, arg_expr, arg_name) -> let deps = arg_param_deps label arg_expr in - if Set.Poly.is_empty deps then None - else Some (cf_loc, deps, arg_name, arg_expr.meta.loc) ) ) - -let list_param_dependant_fundefs_cf (mir : Program.Typed.t) : - (string * Location_span.t * string Set.Poly.t * string * Location_span.t) - Set.Poly.t = + if Set.Poly.is_empty deps + then None + else Some (cf_loc, deps, arg_name, arg_expr.meta.loc))) +;; + +let list_param_dependant_fundefs_cf (mir : Program.Typed.t) + : (string * Location_span.t * string Set.Poly.t * string * Location_span.t) Set.Poly.t + = let info_map = log_prob_build_dep_info_map mir in union_map (Set.Poly.of_list mir.functions_block) ~f:(fun fun_def -> - let dependant_args = - list_param_dependant_fundef_cf mir info_map fun_def - in + let dependant_args = list_param_dependant_fundef_cf mir info_map fun_def in Set.Poly.map dependant_args ~f:(fun (cf_loc, deps, arg_name, arg_loc) -> - (fun_def.fdname, cf_loc, deps, arg_name, arg_loc) ) ) + fun_def.fdname, cf_loc, deps, arg_name, arg_loc)) +;; -let list_non_one_priors (fg : factor_graph) (mir : Program.Typed.t) : - (string * int) Set.Poly.t = +let list_non_one_priors (fg : factor_graph) (mir : Program.Typed.t) + : (string * int) Set.Poly.t + = (* Use the factor graph definition of priors, which treats a neighboring factor as a prior for parameter P if it has no connection to the data except through P *) let priors = list_priors ~factor_graph:(Some fg) mir in let prior_set = - Map.Poly.fold priors ~init:Set.Poly.empty - ~f:(fun ~key:(VVar v) ~data:factors_opt s -> + Map.Poly.fold priors ~init:Set.Poly.empty ~f:(fun ~key:(VVar v) ~data:factors_opt s -> Option.value_map factors_opt ~default:s ~f:(fun factors -> - Set.Poly.add s (v, Set.Poly.length factors) ) ) + Set.Poly.add s (v, Set.Poly.length factors))) in (* Return only multi-prior parameters *) Set.Poly.filter prior_set ~f:(fun (_, n) -> n <> 1) +;; (* Collect useful information about an expression that's available at compile-time into a convenient form. *) let compiletime_value_of_expr (params : (string * Expr.Typed.t Program.transformation) Set.Poly.t) - (data : string Set.Poly.t) (expr : Expr.Typed.t) : - compiletime_val * Expr.Typed.Meta.t = + (data : string Set.Poly.t) + (expr : Expr.Typed.t) + : compiletime_val * Expr.Typed.Meta.t + = let v = match expr with - | {pattern= Var pname; _} -> ( - match Set.Poly.find params ~f:(fun (name, _) -> name = pname) with + | { pattern = Var pname; _ } -> + (match Set.Poly.find params ~f:(fun (name, _) -> name = pname) with | Some (name, trans) -> Param (name, trans) - | None -> ( - match Set.Poly.find data ~f:(fun name -> name = pname) with + | None -> + (match Set.Poly.find data ~f:(fun name -> name = pname) with | Some name -> Data name - | None -> Opaque ) ) + | None -> Opaque)) | _ -> - Option.value_map (num_expr_value expr) ~default:Opaque - ~f:(fun (v, s) -> Number (v, s) ) + Option.value_map (num_expr_value expr) ~default:Opaque ~f:(fun (v, s) -> + Number (v, s)) in - (v, expr.meta) + v, expr.meta +;; (* Scrape all distributions from the program by searching for their function names and function type, and wrangle some useful data about them, like the @@ -255,42 +279,43 @@ let compiletime_value_of_expr let list_distributions (mir : Program.Typed.t) : dist_info Set.Poly.t = let take_dist (expr : Expr.Typed.t) = match expr.pattern with - | Expr.Fixed.Pattern.FunApp (StanLib, fname, arg_exprs) -> ( - match chop_dist_name fname with + | Expr.Fixed.Pattern.FunApp (StanLib, fname, arg_exprs) -> + (match chop_dist_name fname with | Some dname -> - let params = parameter_set mir in - let data = data_set mir in - let args = - List.map ~f:(compiletime_value_of_expr params data) arg_exprs - in - Some {name= dname; loc= expr.meta.loc; args} - | _ -> None ) + let params = parameter_set mir in + let data = data_set mir in + let args = List.map ~f:(compiletime_value_of_expr params data) arg_exprs in + Some { name = dname; loc = expr.meta.loc; args } + | _ -> None) | _ -> None in stmts_collect_exprs - (List.append mir.log_prob - (List.map ~f:(fun f -> f.fdbody) mir.functions_block)) + (List.append mir.log_prob (List.map ~f:(fun f -> f.fdbody) mir.functions_block)) ~f:take_dist +;; (* Our definition of 'unscaled' for constants used in distributions *) let is_unscaled_value (v : float) = let mag = Float.abs v in (mag < 0.1 || mag > 10.0) && mag <> 0.0 +;; -let list_unscaled_constants (distributions_list : dist_info Set.Poly.t) : - (Location_span.t * string) Set.Poly.t = +let list_unscaled_constants (distributions_list : dist_info Set.Poly.t) + : (Location_span.t * string) Set.Poly.t + = (* Search all distributions for unscaled values *) let collect_unscaled_expr (arg : compiletime_val * Expr.Typed.Meta.t) = match arg with | Number (num, num_str), meta -> - if is_unscaled_value num then Set.Poly.singleton (meta.loc, num_str) - else Set.Poly.empty + if is_unscaled_value num + then Set.Poly.singleton (meta.loc, num_str) + else Set.Poly.empty | _ -> Set.Poly.empty in union_map - ~f:(fun {args; _} -> - Set.Poly.union_list (List.map ~f:collect_unscaled_expr args) ) + ~f:(fun { args; _ } -> Set.Poly.union_list (List.map ~f:collect_unscaled_expr args)) distributions_list +;; (********************* Printing functions @@ -300,10 +325,10 @@ let list_unscaled_constants (distributions_list : dist_info Set.Poly.t) : "Warning at :" line and is indented by 2 spaces. *) let pp_warning ppf (loc, msg) = let loc_str = - if loc = Location_span.empty then "" - else " at " ^ Location_span.to_string loc + if loc = Location_span.empty then "" else " at " ^ Location_span.to_string loc in Fmt.pf ppf "Warning%s:@\n@[ %a@]\n" loc_str Fmt.text msg +;; (* Print a set of 'warnings', where each warning comes with its location. By tupling with location and using a set, we're also sorting the warning @@ -317,119 +342,132 @@ let pp_warning ppf (loc, msg) = *) let print_warning_set (warnings : (Location_span.t * string) Set.Poly.t) = let str = - Fmt.strf "%a" - (Fmt.list ~sep:Fmt.nop pp_warning) - (Set.Poly.to_list warnings) + Fmt.strf "%a" (Fmt.list ~sep:Fmt.nop pp_warning) (Set.Poly.to_list warnings) in Out_channel.output_string Out_channel.stderr str +;; let unscaled_constants_message (name : string) : string = Printf.sprintf - "Argument %s suggests there may be parameters that are not unit scale; \ - consider rescaling with a multiplier (see manual section 22.12)." + "Argument %s suggests there may be parameters that are not unit scale; consider \ + rescaling with a multiplier (see manual section 22.12)." name +;; let unscaled_constants_warnings (distributions_list : dist_info Set.Poly.t) = Set.Poly.map - ~f:(fun (loc, name) -> (loc, unscaled_constants_message name)) + ~f:(fun (loc, name) -> loc, unscaled_constants_message name) (list_unscaled_constants distributions_list) +;; let nonsense_constrained_message (pname : string) : string = Printf.sprintf - "Parameter %s has constraints that don't make sense. The lower bound \ - should be strictly less than the upper bound." + "Parameter %s has constraints that don't make sense. The lower bound should be \ + strictly less than the upper bound." pname +;; let hard_constrained_message (pname : string) : string = Printf.sprintf "Your Stan program has a parameter %s with a lower and upper bound in its \ - declaration. These hard constraints are not recommended, for two \ - reasons: (a) Except when there are logical or physical constraints, it \ - is very unusual for you to be sure that a parameter will fall inside a \ - specified range, and (b) The infinite gradient induced by a hard \ - constraint can cause difficulties for Stan's sampling algorithm. As a \ - consequence, we recommend soft constraints rather than hard constraints; \ - for example, instead of constraining an elasticity parameter to fall \ - between 0, and 1, leave it unconstrained and give it a normal(0.5,0.5) \ - prior distribution." + declaration. These hard constraints are not recommended, for two reasons: (a) \ + Except when there are logical or physical constraints, it is very unusual for you \ + to be sure that a parameter will fall inside a specified range, and (b) The \ + infinite gradient induced by a hard constraint can cause difficulties for Stan's \ + sampling algorithm. As a consequence, we recommend soft constraints rather than \ + hard constraints; for example, instead of constraining an elasticity parameter to \ + fall between 0, and 1, leave it unconstrained and give it a normal(0.5,0.5) prior \ + distribution." pname +;; let hard_constrained_warnings (mir : Program.Typed.t) = let pnames = list_hard_constrained mir in Set.Poly.map ~f:(fun (pname, c) -> match c with - | `HardConstraint -> (Location_span.empty, hard_constrained_message pname) - | `NonsenseConstraint -> - (Location_span.empty, nonsense_constrained_message pname) ) + | `HardConstraint -> Location_span.empty, hard_constrained_message pname + | `NonsenseConstraint -> Location_span.empty, nonsense_constrained_message pname) pnames +;; let multi_twiddles_message (vname : string) : string = Printf.sprintf - "The parameter %s is on the left-hand side of more than one twiddle \ - statement." + "The parameter %s is on the left-hand side of more than one twiddle statement." vname +;; let multi_twiddles_warnings (mir : Program.Typed.t) = let twds = list_multi_twiddles mir in Set.Poly.map - ~f:(fun (vname, locs) -> - (Set.Poly.min_elt_exn locs, multi_twiddles_message vname) ) + ~f:(fun (vname, locs) -> Set.Poly.min_elt_exn locs, multi_twiddles_message vname) twds +;; let param_dependant_cf_message (plist : string Set.Poly.t) : string = let plistStr = String.concat ~sep:", " (Set.Poly.to_list plist) in - Printf.sprintf "A control flow statement depends on parameter(s): %s." - plistStr + Printf.sprintf "A control flow statement depends on parameter(s): %s." plistStr +;; let param_dependant_cf_warnings (mir : Program.Typed.t) = let cfs = list_param_dependant_cf mir in - Set.Poly.map - ~f:(fun (loc, plist) -> (loc, param_dependant_cf_message plist)) - cfs - -let param_dependant_fundef_cf_message (fname : string) - (plist : string Set.Poly.t) (arg_name : string) - (callsite : Location_span.t) : string = + Set.Poly.map ~f:(fun (loc, plist) -> loc, param_dependant_cf_message plist) cfs +;; + +let param_dependant_fundef_cf_message + (fname : string) + (plist : string Set.Poly.t) + (arg_name : string) + (callsite : Location_span.t) + : string + = let plistStr = String.concat ~sep:", " (Set.Poly.to_list plist) in Printf.sprintf - "A control flow statement inside function %s depends on argument %s. At \ - %s, the value of %s depends on parameter(s): %s." - fname arg_name + "A control flow statement inside function %s depends on argument %s. At %s, the \ + value of %s depends on parameter(s): %s." + fname + arg_name (Location_span.to_string callsite) - arg_name plistStr + arg_name + plistStr +;; let param_dependant_fundef_cf_warnings (mir : Program.Typed.t) = Set.Poly.map ~f:(fun (fname, cf_loc, deps, arg_name, arg_loc) -> - (cf_loc, param_dependant_fundef_cf_message fname deps arg_name arg_loc) - ) + cf_loc, param_dependant_fundef_cf_message fname deps arg_name arg_loc) (list_param_dependant_fundefs_cf mir) +;; let unused_params_message (pname : string) : string = Printf.sprintf "The parameter %s was declared but was not used in the density calculation." pname +;; -let unused_params_warnings (factor_graph : factor_graph) - (mir : Program.Typed.t) = +let unused_params_warnings (factor_graph : factor_graph) (mir : Program.Typed.t) = Set.Poly.map - ~f:(fun pname -> (Location_span.empty, unused_params_message pname)) + ~f:(fun pname -> Location_span.empty, unused_params_message pname) (list_unused_params factor_graph mir) +;; let non_one_priors_message (pname : string) (n : int) : string = - if n = 0 then Printf.sprintf "The parameter %s has no priors." pname + if n = 0 + then Printf.sprintf "The parameter %s has no priors." pname else Printf.sprintf "The parameter %s has %d priors." pname n +;; -let non_one_priors_warnings (factor_graph : factor_graph) - (mir : Program.Typed.t) = +let non_one_priors_warnings (factor_graph : factor_graph) (mir : Program.Typed.t) = Set.Poly.map - ~f:(fun (pname, n) -> (Location_span.empty, non_one_priors_message pname n)) + ~f:(fun (pname, n) -> Location_span.empty, non_one_priors_message pname n) (list_non_one_priors factor_graph mir) +;; let uninitialized_message (vname : string) : string = Printf.sprintf - "The variable %s may not have been assigned a value before its use." vname + "The variable %s may not have been assigned a value before its use." + vname +;; let uninitialized_warnings (mir : Program.Typed.t) = let uninit_vars = @@ -437,28 +475,26 @@ let uninitialized_warnings (mir : Program.Typed.t) = ~f:(fun (span, _) -> span <> Location_span.empty) (Dependence_analysis.mir_uninitialized_variables mir) in - Set.Poly.map - ~f:(fun (loc, vname) -> (loc, uninitialized_message vname)) - uninit_vars + Set.Poly.map ~f:(fun (loc, vname) -> loc, uninitialized_message vname) uninit_vars +;; (* Print uninitialized warnings In case a user wants only this warning *) -let print_warn_uninitialized mir = - print_warning_set (uninitialized_warnings mir) +let print_warn_uninitialized mir = print_warning_set (uninitialized_warnings mir) (* Optimization settings for constant propagation and partial evaluation *) let settings_constant_prop = { no_optimizations with - constant_propagation= true - ; copy_propagation= true - ; partial_evaluation= true } + constant_propagation = true + ; copy_propagation = true + ; partial_evaluation = true + } +;; (* Print all pedantic mode warnings, sorted, to stderr *) let print_warn_pedantic (mir_unopt : Program.Typed.t) = (* Some warnings will be stronger when constants are propagated *) - let mir = - Optimize.optimization_suite ~settings:settings_constant_prop mir_unopt - in + let mir = Optimize.optimization_suite ~settings:settings_constant_prop mir_unopt in (* Try to avoid recomputation by pre-building structures *) let distributions_info = list_distributions mir in let factor_graph = prog_factor_graph mir in @@ -472,6 +508,8 @@ let print_warn_pedantic (mir_unopt : Program.Typed.t) = ; param_dependant_cf_warnings mir ; param_dependant_fundef_cf_warnings mir ; non_one_priors_warnings factor_graph mir - ; distribution_warnings distributions_info ] + ; distribution_warnings distributions_info + ] in print_warning_set warning_set +;; diff --git a/src/analysis_and_optimization/Pedantic_analysis.mli b/src/analysis_and_optimization/Pedantic_analysis.mli index 9ab25f1b10..03d70c8504 100644 --- a/src/analysis_and_optimization/Pedantic_analysis.mli +++ b/src/analysis_and_optimization/Pedantic_analysis.mli @@ -1,11 +1,11 @@ open Middle -val print_warn_pedantic : Program.Typed.t -> unit (** Print all pedantic mode warnings to stderr. *) +val print_warn_pedantic : Program.Typed.t -> unit -val print_warn_uninitialized : Program.Typed.t -> unit (** Print warnings about each variable which is used before being initialized *) +val print_warn_uninitialized : Program.Typed.t -> unit diff --git a/src/analysis_and_optimization/Pedantic_dist_warnings.ml b/src/analysis_and_optimization/Pedantic_dist_warnings.ml index 0b5de8cad0..1b6f69e374 100644 --- a/src/analysis_and_optimization/Pedantic_dist_warnings.ml +++ b/src/analysis_and_optimization/Pedantic_dist_warnings.ml @@ -17,12 +17,16 @@ type compiletime_val = distribution properties are met *) type dist_info = - { name: string - ; loc: Location_span.t - ; args: (compiletime_val * Expr.Typed.Meta.t) List.t } + { name : string + ; loc : Location_span.t + ; args : (compiletime_val * Expr.Typed.Meta.t) List.t + } (* Value constraint as a range. The bools are true if the bound is inclusive *) -type range = {lower: (float * bool) option; upper: (float * bool) option} +type range = + { lower : (float * bool) option + ; upper : (float * bool) option + } (* Value constraint for an argument *) type var_constraint = @@ -37,49 +41,60 @@ type var_constraint = | Covariance (* Constraint paired with a name for user messages *) -type var_constraint_named = {name: string; constr: var_constraint} +type var_constraint_named = + { name : string + ; constr : var_constraint + } let unit_range = - { name= "[0,1]" - ; constr= Range {lower= Some (0., true); upper= Some (1., true)} } + { name = "[0,1]"; constr = Range { lower = Some (0., true); upper = Some (1., true) } } +;; let exclusive_unit_range = - { name= "(0,1)" - ; constr= Range {lower= Some (0., false); upper= Some (1., false)} } + { name = "(0,1)" + ; constr = Range { lower = Some (0., false); upper = Some (1., false) } + } +;; let positive_range = - { name= "strictly positive" - ; constr= Range {lower= Some (0., false); upper= None} } + { name = "strictly positive" + ; constr = Range { lower = Some (0., false); upper = None } + } +;; let nonnegative_range = - {name= "non-negative"; constr= Range {lower= Some (0., true); upper= None}} + { name = "non-negative"; constr = Range { lower = Some (0., true); upper = None } } +;; -let simplex = {name= "simplex"; constr= Simplex} -let ordered = {name= "ordered"; constr= Ordered} -let correlation = {name= "correlation"; constr= Correlation} +let simplex = { name = "simplex"; constr = Simplex } +let ordered = { name = "ordered"; constr = Ordered } +let correlation = { name = "correlation"; constr = Correlation } let cholesky_correlation = - {name= "Cholesky factor of correlation"; constr= CholeskyCorr} + { name = "Cholesky factor of correlation"; constr = CholeskyCorr } +;; -let covariance = {name= "covariance"; constr= Covariance} - -let cholesky_covariance = - {name= "Cholesky factor of covariance"; constr= CholeskyCov} +let covariance = { name = "covariance"; constr = Covariance } +let cholesky_covariance = { name = "Cholesky factor of covariance"; constr = CholeskyCov } (* Check for inconsistency between a distribution argument's value range and the declared bounds of a variable *) let bounds_out_of_range (range : range) (bounds : bound_values) : bool = - match (bounds.lower, bounds.upper, range.lower, range.upper) with + match bounds.lower, bounds.upper, range.lower, range.upper with | `None, _, Some _, _ -> true | _, `None, _, Some _ -> true | `Lit l, _, Some (l', _), _ when l < l' -> true | _, `Lit u, _, Some (u', _) when u > u' -> true | _ -> false +;; (* Check for inconsistency between a distribution argument's constraint and the constraint transformation of a variable *) -let transform_mismatch_constraint (constr : var_constraint) - (trans : Expr.Typed.t Program.transformation) : bool = +let transform_mismatch_constraint + (constr : var_constraint) + (trans : Expr.Typed.t Program.transformation) + : bool + = match constr with | Range range -> bounds_out_of_range range (trans_bounds_values trans) | Ordered -> trans <> Program.Ordered @@ -87,10 +102,10 @@ let transform_mismatch_constraint (constr : var_constraint) | Simplex -> trans <> Program.Simplex | UnitVector -> trans <> Program.UnitVector | CholeskyCorr -> trans <> Program.CholeskyCorr - | CholeskyCov -> - trans <> Program.CholeskyCov && trans <> Program.CholeskyCorr + | CholeskyCov -> trans <> Program.CholeskyCov && trans <> Program.CholeskyCorr | Correlation -> trans <> Program.Correlation | Covariance -> trans <> Program.Covariance && trans <> Program.Correlation +;; (* Check for inconsistency between a distribution argument's range and a literal value *) @@ -108,6 +123,7 @@ let value_out_of_range (range : range) (v : float) = | None -> false in lower_bad || upper_bad +;; (* Check for inconsistency between a distribution argument's constraint and a literal value *) @@ -117,71 +133,110 @@ let value_mismatch_constraint (constr : var_constraint) (v : float) = (* We don't know how to check if a value falls into a constraint other than Range, unless we want to inspect e.g. matrix literals. *) | _ -> false +;; (********************* Argument constraint mismatch warnings ********************) -type arg_info = Arg of (int * string) | Variate +type arg_info = + | Arg of (int * string) + | Variate let arg_number (arg : arg_info) = - match arg with Arg (n, _) -> n | Variate -> 0 - -let constr_mismatch_message (dist_name : string) (param_name : string) - (arg : arg_info) (constr_name : string) : string = + match arg with + | Arg (n, _) -> n + | Variate -> 0 +;; + +let constr_mismatch_message + (dist_name : string) + (param_name : string) + (arg : arg_info) + (constr_name : string) + : string + = match arg with | Arg (argn, arg_name) -> - Printf.sprintf - "A %s distribution is given parameter %s as %s (argument %d), but %s \ - was not constrained to be %s." - dist_name param_name arg_name argn param_name constr_name + Printf.sprintf + "A %s distribution is given parameter %s as %s (argument %d), but %s was not \ + constrained to be %s." + dist_name + param_name + arg_name + argn + param_name + constr_name | Variate -> - (* Possibly: Either change the distribution or change the constraints. *) - Printf.sprintf - "Parameter %s is given a %s distribution, which has %s support, but \ - %s was not constrained to be %s." - param_name dist_name constr_name param_name constr_name - -let constr_literal_mismatch_message (dist_name : string) (num_str : string) - (arg : arg_info) (constr_name : string) : string = + (* Possibly: Either change the distribution or change the constraints. *) + Printf.sprintf + "Parameter %s is given a %s distribution, which has %s support, but %s was not \ + constrained to be %s." + param_name + dist_name + constr_name + param_name + constr_name +;; + +let constr_literal_mismatch_message + (dist_name : string) + (num_str : string) + (arg : arg_info) + (constr_name : string) + : string + = match arg with | Arg (argn, arg_name) -> - Printf.sprintf - "A %s distribution is given value %s as %s (argument %d), but %s is \ - not %s." - dist_name num_str arg_name argn arg_name constr_name + Printf.sprintf + "A %s distribution is given value %s as %s (argument %d), but %s is not %s." + dist_name + num_str + arg_name + argn + arg_name + constr_name | Variate -> - (* Possibly: Either change the distribution or change the constraints. *) - Printf.sprintf - "Value %s is given a %s distribution, which has %s support, but %s is \ - not %s." - num_str dist_name constr_name num_str constr_name + (* Possibly: Either change the distribution or change the constraints. *) + Printf.sprintf + "Value %s is given a %s distribution, which has %s support, but %s is not %s." + num_str + dist_name + constr_name + num_str + constr_name +;; (* Return a warning if the argn-th argument doesn't match its constraints *) -let constr_mismatch_warning (constr : var_constraint_named) (arg : arg_info) - ({args; name; loc} : dist_info) : (Location_span.t * string) option = +let constr_mismatch_warning + (constr : var_constraint_named) + (arg : arg_info) + ({ args; name; loc } : dist_info) + : (Location_span.t * string) option + = let v = match List.nth args (arg_number arg) with | Some v -> v | None -> - let arg_fail_msg = - Printf.sprintf "Distribution %s at %s expects more arguments." name - (Location_span.to_string loc) - in - raise (Failure arg_fail_msg) + let arg_fail_msg = + Printf.sprintf + "Distribution %s at %s expects more arguments." + name + (Location_span.to_string loc) + in + raise (Failure arg_fail_msg) in match v with | Param (pname, trans), meta -> - if transform_mismatch_constraint constr.constr trans then - Some (meta.loc, constr_mismatch_message name pname arg constr.name) - else None + if transform_mismatch_constraint constr.constr trans + then Some (meta.loc, constr_mismatch_message name pname arg constr.name) + else None | Number (num, num_str), meta -> - if value_mismatch_constraint constr.constr num then - Some - ( meta.loc - , constr_literal_mismatch_message name num_str arg constr.name ) - else None + if value_mismatch_constraint constr.constr num + then Some (meta.loc, constr_literal_mismatch_message name num_str arg constr.name) + else None | _ -> None +;; (********************* Distribution-specific warnings @@ -189,67 +244,67 @@ let constr_mismatch_warning (constr : var_constraint_named) (arg : arg_info) let uniform_dist_message (pname : string) : string = Printf.sprintf - "Parameter %s is given a uniform distribution. The uniform distribution \ - is not recommended, for two reasons: (a) Except when there are logical \ - or physical constraints, it is very unusual for you to be sure that a \ - parameter will fall inside a specified range, and (b) The infinite \ - gradient induced by a uniform density can cause difficulties for Stan's \ - sampling algorithm. As a consequence, we recommend soft constraints \ - rather than hard constraints; for example, instead of giving an \ - elasticity parameter a uniform(0,1) distribution, try normal(0.5,0.5)." + "Parameter %s is given a uniform distribution. The uniform distribution is not \ + recommended, for two reasons: (a) Except when there are logical or physical \ + constraints, it is very unusual for you to be sure that a parameter will fall \ + inside a specified range, and (b) The infinite gradient induced by a uniform \ + density can cause difficulties for Stan's sampling algorithm. As a consequence, we \ + recommend soft constraints rather than hard constraints; for example, instead of \ + giving an elasticity parameter a uniform(0,1) distribution, try normal(0.5,0.5)." pname +;; (* Warning for all uniform distributions with a parameter *) -let uniform_dist_warning (dist_info : dist_info) : - (Location_span.t * string) option = +let uniform_dist_warning (dist_info : dist_info) : (Location_span.t * string) option = match dist_info with - | {args= (Param (pname, trans), _) :: (arg1, _) :: (arg2, _) :: _; _} -> ( - let warning = Some (dist_info.loc, uniform_dist_message pname) in - match (arg1, arg2, trans_bounds_values trans) with - | _, _, {upper= `None; _} | _, _, {lower= `None; _} -> - (* the variate is unbounded *) - warning - | Number (uni, _), _, {lower= `Lit bound; _} - |_, Number (uni, _), {upper= `Lit bound; _} -> - (* the variate is bounded differently than the uniform dist *) - if uni = bound then None else warning - | _ -> None ) + | { args = (Param (pname, trans), _) :: (arg1, _) :: (arg2, _) :: _; _ } -> + let warning = Some (dist_info.loc, uniform_dist_message pname) in + (match arg1, arg2, trans_bounds_values trans with + | _, _, { upper = `None; _ } | _, _, { lower = `None; _ } -> + (* the variate is unbounded *) + warning + | Number (uni, _), _, { lower = `Lit bound; _ } + | _, Number (uni, _), { upper = `Lit bound; _ } -> + (* the variate is bounded differently than the uniform dist *) + if uni = bound then None else warning + | _ -> None) | _ -> None +;; let lkj_corr_message : string = "It is suggested to reparameterize your model to replace lkj_corr with \ - lkj_corr_cholesky, the Cholesky factor variant. lkj_corr tends to run \ - slower, consume more memory, and has higher risk of numerical errors." + lkj_corr_cholesky, the Cholesky factor variant. lkj_corr tends to run slower, consume \ + more memory, and has higher risk of numerical errors." +;; (* Warn about all non-Cholesky lkj_corr distributions *) -let lkj_corr_dist_warning (dist_info : dist_info) : - (Location_span.t * string) option = +let lkj_corr_dist_warning (dist_info : dist_info) : (Location_span.t * string) option = Some (dist_info.loc, lkj_corr_message) +;; let gamma_arg_dist_message : string = - "There is a gamma or inverse-gamma distribution with parameters that are \ - equal to each other and set to values less than 1. This is mathematically \ - acceptable and can make sense in some problems, but typically we see this \ - model used as an attempt to assign a noninformative prior distribution. In \ - fact, priors such as inverse-gamma(.001,.001) can be very strong, as \ - explained by Gelman (2006). Instead we recommend something like a \ - normal(0,1) or student_t(4,0,1), with parameter constrained to be positive." + "There is a gamma or inverse-gamma distribution with parameters that are equal to each \ + other and set to values less than 1. This is mathematically acceptable and can make \ + sense in some problems, but typically we see this model used as an attempt to assign \ + a noninformative prior distribution. In fact, priors such as inverse-gamma(.001,.001) \ + can be very strong, as explained by Gelman (2006). Instead we recommend something \ + like a normal(0,1) or student_t(4,0,1), with parameter constrained to be positive." +;; (* Warning particular to gamma and inv_gamma, when A=B<1 *) -let gamma_arg_dist_warning (dist_info : dist_info) : - (Location_span.t * string) option = +let gamma_arg_dist_warning (dist_info : dist_info) : (Location_span.t * string) option = match dist_info with - | {args= [_; (Number (a, _), meta); (Number (b, _), _)]; _} -> - if a = b && a < 1. then Some (meta.loc, gamma_arg_dist_message) else None + | { args = [ _; (Number (a, _), meta); (Number (b, _), _) ]; _ } -> + if a = b && a < 1. then Some (meta.loc, gamma_arg_dist_message) else None | _ -> None +;; (********************* Distribution properties table ********************) (* Generate all of the warnings that are relevant to a given distribution *) -let distribution_warning (dist_info : dist_info) : - (Location_span.t * string) List.t = +let distribution_warning (dist_info : dist_info) : (Location_span.t * string) List.t = let scale_name = "a scale parameter" in let scale_mat_name = "a scale matrix" in let inv_scale_name = "an inverse scale parameter" in @@ -263,214 +318,241 @@ let distribution_warning (dist_info : dist_info) : match dist_info.name with (* Binary Distributions *) | "bernoulli" -> - [ (* Note: variate binary *) - constr_mismatch_warning unit_range (Arg (1, "chance of success")) ] + [ (* Note: variate binary *) + constr_mismatch_warning unit_range (Arg (1, "chance of success")) + ] | "bernoulli_logit" -> [ (* Note: variate binary *) ] | "bernoulli_logit_glm" -> [ (* Note: variate binary *) ] (* Bounded Discrete Distributions *) | "binomial" -> - [ (* Note: variate nonnegative int *) - (* Note: args 1 nonnegative int *) - constr_mismatch_warning unit_range (Arg (2, "chance of success")) ] + [ (* Note: variate nonnegative int *) + (* Note: args 1 nonnegative int *) + constr_mismatch_warning unit_range (Arg (2, "chance of success")) + ] | "binomial_logit" -> - [ (* Note: variate nonnegative int *) - (* Note: args 1 nonnegative int *) ] + [ (* Note: variate nonnegative int *) + (* Note: args 1 nonnegative int *) ] | "beta_binomial" -> - [ (* Note: variate nonnegative int *) - (* Note: args 1 nonnegative int *) - constr_mismatch_warning positive_range - (Arg (2, "a prior success count")) - ; constr_mismatch_warning positive_range - (Arg (3, "a prior failure count")) ] + [ (* Note: variate nonnegative int *) + (* Note: args 1 nonnegative int *) + constr_mismatch_warning positive_range (Arg (2, "a prior success count")) + ; constr_mismatch_warning positive_range (Arg (3, "a prior failure count")) + ] | "hypergeometric" -> - [ (* Note: variate nonnegative int *) - (* Note: args 1,2,3 nonnegative int *) ] + [ (* Note: variate nonnegative int *) + (* Note: args 1,2,3 nonnegative int *) ] | "categorical" -> - [ (* Note: variate positive int *) - constr_mismatch_warning simplex - (Arg (1, "a vector of outcome probabilities")) ] + [ (* Note: variate positive int *) + constr_mismatch_warning simplex (Arg (1, "a vector of outcome probabilities")) + ] | "ordered_logistic" -> - [ (* Note: variate positive int *) - constr_mismatch_warning ordered (Arg (2, "cutpoints")) ] + [ (* Note: variate positive int *) + constr_mismatch_warning ordered (Arg (2, "cutpoints")) + ] | "ordered_probit" -> - [ (* Note: variate positive int *) - constr_mismatch_warning ordered (Arg (2, "cutpoints")) ] + [ (* Note: variate positive int *) + constr_mismatch_warning ordered (Arg (2, "cutpoints")) + ] (* Unbounded Discrete Distributions *) | "neg_binomial" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range (Arg (2, inv_scale_name)) ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, inv_scale_name)) + ] | "neg_binomial_2" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range - (Arg (2, "a precision parameter")) ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, "a precision parameter")) + ] | "neg_binomial_2_log" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning positive_range - (Arg (2, "an inverse overdispersion control parameter")) ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning + positive_range + (Arg (2, "an inverse overdispersion control parameter")) + ] | "neg_binomial_2_log_glm" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning positive_range - (Arg (4, "an inverse overdispersion control parameter")) ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning + positive_range + (Arg (4, "an inverse overdispersion control parameter")) + ] | "poisson" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning positive_range (Arg (1, "a rate parameter")) - ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning positive_range (Arg (1, "a rate parameter")) + ] | "poisson_log" -> [ (* Note: variate nonnegative int *) ] | "poisson_log_glm" -> [ (* Note: variate nonnegative int *) ] (* Multivariate Discrete Distributions *) | "multinomial" -> - [ (* Note: variate nonnegative int *) - constr_mismatch_warning simplex (Arg (1, "a distribution parameter")) - ] + [ (* Note: variate nonnegative int *) + constr_mismatch_warning simplex (Arg (1, "a distribution parameter")) + ] (* Unbounded Continuous Distributions *) - | "normal" -> [constr_mismatch_warning positive_range (Arg (2, scale_name))] - | "normal_id_glm" -> - [constr_mismatch_warning positive_range (Arg (4, scale_name))] + | "normal" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + | "normal_id_glm" -> [ constr_mismatch_warning positive_range (Arg (4, scale_name)) ] | "exp_mod_normal" -> - [ constr_mismatch_warning positive_range (Arg (2, scale_name)) - ; constr_mismatch_warning positive_range (Arg (3, shape_name)) ] - | "skew_normal" -> - [constr_mismatch_warning positive_range (Arg (2, scale_name))] + [ constr_mismatch_warning positive_range (Arg (2, scale_name)) + ; constr_mismatch_warning positive_range (Arg (3, shape_name)) + ] + | "skew_normal" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] | "student_t" -> - [ constr_mismatch_warning positive_range (Arg (1, dof_name)) - ; constr_mismatch_warning positive_range (Arg (3, scale_name)) ] - | "cauchy" -> [constr_mismatch_warning positive_range (Arg (2, scale_name))] + [ constr_mismatch_warning positive_range (Arg (1, dof_name)) + ; constr_mismatch_warning positive_range (Arg (3, scale_name)) + ] + | "cauchy" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] | "double_exponential" -> - [constr_mismatch_warning positive_range (Arg (2, scale_name))] - | "logistic" -> - [constr_mismatch_warning positive_range (Arg (2, scale_name))] - | "gumbel" -> [constr_mismatch_warning positive_range (Arg (2, scale_name))] + [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + | "logistic" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + | "gumbel" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] (* Positive Continuous Distributions *) | "lognormal" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (2, scale_name)) + ] | "chi_square" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, dof_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, dof_name)) + ] | "inv_chi_square" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, dof_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, dof_name)) + ] | "scaled_inv_chi_square" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, dof_name)) - ; constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, dof_name)) + ; constr_mismatch_warning positive_range (Arg (2, scale_name)) + ] | "exponential" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, scale_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, scale_name)) + ] | "gamma" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range (Arg (2, inv_scale_name)) - ; gamma_arg_dist_warning ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, inv_scale_name)) + ; gamma_arg_dist_warning + ] | "inv_gamma" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range (Arg (2, scale_name)) - ; gamma_arg_dist_warning ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, scale_name)) + ; gamma_arg_dist_warning + ] | "weibull" -> - [ constr_mismatch_warning nonnegative_range Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + [ constr_mismatch_warning nonnegative_range Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, scale_name)) + ] | "frechet" -> - [ constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) - ; constr_mismatch_warning positive_range (Arg (2, scale_name)) ] + [ constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ; constr_mismatch_warning positive_range (Arg (2, scale_name)) + ] (* Non-negative Continuous Distributions *) | "rayleigh" -> - [ constr_mismatch_warning nonnegative_range Variate - ; constr_mismatch_warning positive_range (Arg (1, scale_name)) ] + [ constr_mismatch_warning nonnegative_range Variate + ; constr_mismatch_warning positive_range (Arg (1, scale_name)) + ] | "wiener" -> - [ (* Note: Could do more here, since variate should be > arg 2 *) - constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range - (Arg (1, "a boundary separation parameter")) - ; constr_mismatch_warning positive_range - (Arg (2, "a non-decision time parameter")) - ; constr_mismatch_warning unit_range - (Arg (3, "an a-priori bias parameter")) ] + [ (* Note: Could do more here, since variate should be > arg 2 *) + constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning + positive_range + (Arg (1, "a boundary separation parameter")) + ; constr_mismatch_warning positive_range (Arg (2, "a non-decision time parameter")) + ; constr_mismatch_warning unit_range (Arg (3, "an a-priori bias parameter")) + ] (* Positive Lower-Bounded Probabilities *) | "pareto" -> - [ (* Note: Variate >= arg 1 *) - constr_mismatch_warning positive_range Variate - ; constr_mismatch_warning positive_range - (Arg (1, "a positive minimum parameter")) - ; constr_mismatch_warning positive_range (Arg (2, shape_name)) ] + [ (* Note: Variate >= arg 1 *) + constr_mismatch_warning positive_range Variate + ; constr_mismatch_warning positive_range (Arg (1, "a positive minimum parameter")) + ; constr_mismatch_warning positive_range (Arg (2, shape_name)) + ] | "pareto_type_2" -> - [ (* Note: Variate >= arg 1 *) - constr_mismatch_warning positive_range (Arg (2, scale_name)) - ; constr_mismatch_warning positive_range (Arg (3, shape_name)) ] + [ (* Note: Variate >= arg 1 *) + constr_mismatch_warning positive_range (Arg (2, scale_name)) + ; constr_mismatch_warning positive_range (Arg (3, shape_name)) + ] (* Continuous Distributions on [0,1] *) | "beta" -> - [ constr_mismatch_warning exclusive_unit_range Variate - ; constr_mismatch_warning positive_range (Arg (1, "a count parameter")) - ; constr_mismatch_warning positive_range (Arg (2, "a count parameter")) - ] + [ constr_mismatch_warning exclusive_unit_range Variate + ; constr_mismatch_warning positive_range (Arg (1, "a count parameter")) + ; constr_mismatch_warning positive_range (Arg (2, "a count parameter")) + ] | "beta_proportion" -> - [ constr_mismatch_warning exclusive_unit_range Variate - ; constr_mismatch_warning exclusive_unit_range - (Arg (1, "a unit mean parameter")) - ; constr_mismatch_warning positive_range - (Arg (2, "a precision parameter")) ] + [ constr_mismatch_warning exclusive_unit_range Variate + ; constr_mismatch_warning exclusive_unit_range (Arg (1, "a unit mean parameter")) + ; constr_mismatch_warning positive_range (Arg (2, "a precision parameter")) + ] (* Circular Distributions *) - | "von_mises" -> - [constr_mismatch_warning positive_range (Arg (2, scale_name))] + | "von_mises" -> [ constr_mismatch_warning positive_range (Arg (2, scale_name)) ] (* Bounded Continuous Distributions *) | "uniform" -> - [ (* Could also check b > c *) - (* Can this be generalized, by restricting a < variate < b? *) - uniform_dist_warning ] + [ (* Could also check b > c *) + (* Can this be generalized, by restricting a < variate < b? *) + uniform_dist_warning + ] (* Distributions over Unbounded Vectors *) - | "multi_normal" -> [constr_mismatch_warning covariance (Arg (2, cov_name))] + | "multi_normal" -> [ constr_mismatch_warning covariance (Arg (2, cov_name)) ] | "multi_normal_prec" -> - [constr_mismatch_warning covariance (Arg (2, "a precision matrix"))] + [ constr_mismatch_warning covariance (Arg (2, "a precision matrix")) ] | "multi_normal_cholesky" -> - [constr_mismatch_warning cholesky_covariance (Arg (2, cov_name))] + [ constr_mismatch_warning cholesky_covariance (Arg (2, cov_name)) ] | "multi_gp" -> - [ (* Note: arg 2 "inverse scales" is vector of positive inverse scales*) - constr_mismatch_warning covariance (Arg (1, "a kernel matrix")) ] + [ (* Note: arg 2 "inverse scales" is vector of positive inverse scales*) + constr_mismatch_warning covariance (Arg (1, "a kernel matrix")) + ] | "multi_gp_cholesky" -> - [ (* Note: arg 2 "inverse scales" is vector of positive inverse scales*) - constr_mismatch_warning cholesky_covariance - (Arg (1, "Cholesky factor of the kernel matrix")) ] + [ (* Note: arg 2 "inverse scales" is vector of positive inverse scales*) + constr_mismatch_warning + cholesky_covariance + (Arg (1, "Cholesky factor of the kernel matrix")) + ] | "multi_student_t" -> - [ constr_mismatch_warning positive_range (Arg (1, dof_name)) - ; constr_mismatch_warning covariance (Arg (3, scale_mat_name)) ] + [ constr_mismatch_warning positive_range (Arg (1, dof_name)) + ; constr_mismatch_warning covariance (Arg (3, scale_mat_name)) + ] | "gaussian_dlm_obs" -> - [ constr_mismatch_warning covariance - (Arg (3, "observation covariance matrix")) - ; constr_mismatch_warning covariance - (Arg (4, "system covariance matrix")) ] + [ constr_mismatch_warning covariance (Arg (3, "observation covariance matrix")) + ; constr_mismatch_warning covariance (Arg (4, "system covariance matrix")) + ] (* Simplex Distributions *) | "dirichlet" -> - [ constr_mismatch_warning simplex Variate - ; constr_mismatch_warning positive_range (Arg (1, "a count parameter")) - ] + [ constr_mismatch_warning simplex Variate + ; constr_mismatch_warning positive_range (Arg (1, "a count parameter")) + ] (* Correlation Matrix Distributions *) | "lkj_corr" -> - [ lkj_corr_dist_warning - ; constr_mismatch_warning correlation Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) ] + [ lkj_corr_dist_warning + ; constr_mismatch_warning correlation Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ] | "lkj_corr_cholesky" -> - [ constr_mismatch_warning cholesky_correlation Variate - ; constr_mismatch_warning positive_range (Arg (1, shape_name)) ] + [ constr_mismatch_warning cholesky_correlation Variate + ; constr_mismatch_warning positive_range (Arg (1, shape_name)) + ] (* Covariance Matrix Distributions *) | "wishart" -> - [ constr_mismatch_warning covariance Variate - ; constr_mismatch_warning positive_range (Arg (1, dof_name)) - ; constr_mismatch_warning covariance (Arg (2, scale_mat_name)) ] + [ constr_mismatch_warning covariance Variate + ; constr_mismatch_warning positive_range (Arg (1, dof_name)) + ; constr_mismatch_warning covariance (Arg (2, scale_mat_name)) + ] | "inv_wishart" -> - [ constr_mismatch_warning covariance Variate - ; constr_mismatch_warning positive_range (Arg (1, dof_name)) - ; constr_mismatch_warning covariance (Arg (2, scale_mat_name)) ] + [ constr_mismatch_warning covariance Variate + ; constr_mismatch_warning positive_range (Arg (1, dof_name)) + ; constr_mismatch_warning covariance (Arg (2, scale_mat_name)) + ] | _ -> [] in List.filter_map ~f:(fun f -> f dist_info) warning_fns +;; (* Generate the distribution warnings for a program *) -let distribution_warnings (distributions_list : dist_info Set.Poly.t) : - (Location_span.t * string) Set.Poly.t = +let distribution_warnings (distributions_list : dist_info Set.Poly.t) + : (Location_span.t * string) Set.Poly.t + = union_map ~f:(fun dist_info -> Set.Poly.of_list (distribution_warning dist_info)) distributions_list +;; diff --git a/src/common/Fixed.ml b/src/common/Fixed.ml index 5eff911437..d8136434c7 100644 --- a/src/common/Fixed.ml +++ b/src/common/Fixed.ml @@ -7,13 +7,15 @@ open Core_kernel module type S = sig module Pattern : Pattern.S - type 'a t = {pattern: 'a t Pattern.t; meta: 'a} + type 'a t = + { pattern : 'a t Pattern.t + ; meta : 'a + } [@@deriving compare, map, fold, hash, sexp] include Foldable.S with type 'a t := 'a t include Pretty.S1 with type 'a t := 'a t - val fold_pattern : f:('a * 'r Pattern.t -> 'r) -> 'a t -> 'r (** `fold_pattern` traverses the data structure from the bottom up replacing the original meta data of type `'a` with some result type `'r` @@ -41,14 +43,14 @@ module type S = sig meta data `'a`. That is, a tuple `('a * 'b t Pattern.t)` has two type variables where as `'a t` has one. *) + val fold_pattern : f:('a * 'r Pattern.t -> 'r) -> 'a t -> 'r - val rewrite_bottom_up : f:('a t -> 'a t) -> 'a t -> 'a t (** `rewrite_bottom_up` specializes `fold_pattern` so that the result type `'r` is equal to the type of our fixed-point data structure i.e. `'r = 'a t`. This also means that the function `f` can be written with our fixed-point type `'a t` as its argument. *) + val rewrite_bottom_up : f:('a t -> 'a t) -> 'a t -> 'a t - val unfold_pattern : f:('r -> 'a * 'r Pattern.t) -> 'r -> 'a t (** `unfold` builds a fixed-point data structure from the top down. Starting with a user-supplied seed of type `'r`, `unfold` recursively applies the function `f` yielding a tuple of meta-data and pattern with elements of @@ -59,49 +61,59 @@ module type S = sig As with `fold_pattern` the function `f` returns a tuple of meta-data and pattern rather than our record type since, in general, `'r =/= 'a t`. *) + val unfold_pattern : f:('r -> 'a * 'r Pattern.t) -> 'r -> 'a t - val rewrite_top_down : f:('a t -> 'a t) -> 'a t -> 'a t (** `rewrite_top_down` specializes `unfold` by requiring that `'r = 'a t`. As a consequence the function `f` accepts our record type `'a t` as its argument. *) + val rewrite_top_down : f:('a t -> 'a t) -> 'a t -> 'a t end (** Functor which creates the fixed-point of the type defined in the `Pattern` module argument *) module Make (Pattern : Pattern.S) : S with module Pattern := Pattern = struct - type 'a t = {pattern: 'a t Pattern.t; meta: 'a} + type 'a t = + { pattern : 'a t Pattern.t + ; meta : 'a + } [@@deriving compare, map, fold, hash, sexp] - let rec pp f ppf {pattern; meta} = + let rec pp f ppf { pattern; meta } = Fmt.pf ppf {|%a%a|} f meta (Pattern.pp (pp f)) pattern + ;; - include Foldable.Make (struct type nonrec 'a t = 'a t + include Foldable.Make (struct + type nonrec 'a t = 'a t - let fold = fold + let fold = fold end) - let rec fold_pattern ~f {meta; pattern} = + let rec fold_pattern ~f { meta; pattern } = let pattern' = Pattern.map (fold_pattern ~f) pattern in f (meta, pattern') + ;; (** For clarity this is written explicitly but is equivalent to `fold_pattern ~f:(Fn.compose f fix) t` *) let rec rewrite_bottom_up ~f t = - let x = {t with pattern= Pattern.map (rewrite_bottom_up ~f) t.pattern} in + let x = { t with pattern = Pattern.map (rewrite_bottom_up ~f) t.pattern } in f x + ;; let rec unfold_pattern ~f x = let meta, pattern = f x in - {meta; pattern= Pattern.map (unfold_pattern ~f) pattern} + { meta; pattern = Pattern.map (unfold_pattern ~f) pattern } + ;; (** For clarity this is written explicitly but it is equivalent to `unfold ~f:(Fn.compose unfix f) x` *) let rec rewrite_top_down ~f x = let t = f x in - {t with pattern= Pattern.map (rewrite_top_down ~f) t.pattern} + { t with pattern = Pattern.map (rewrite_top_down ~f) t.pattern } + ;; end (** Nested fixed-point type where an element of the `Pattern` is itself @@ -112,29 +124,27 @@ module type S2 = sig module First : S module Pattern : Pattern.S2 - type ('a, 'b) t = {pattern: ('a First.t, ('a, 'b) t) Pattern.t; meta: 'b} + type ('a, 'b) t = + { pattern : ('a First.t, ('a, 'b) t) Pattern.t + ; meta : 'b + } [@@deriving compare, map, fold, hash, sexp] include Foldable.S2 with type ('a, 'b) t := ('a, 'b) t include Pretty.S2 with type ('a, 'b) t := ('a, 'b) t - val fold_pattern : - f:('a * 'r1 First.Pattern.t -> 'r1) - -> g:('b * ('r1, 'r2) Pattern.t -> 'r2) - -> ('a, 'b) t - -> 'r2 (** `fold_pattern` traverses the data structure from the bottom up replacing the original meta data of `First` with type `'a` with some result type `'r1` and the meta data at this level withtype `'b` to another result type `'r2` and combines the result values on the way back up. *) - - val rewrite_bottom_up : - f:('a First.t -> 'a First.t) - -> g:(('a, 'b) t -> ('a, 'b) t) - -> ('a, 'b) t + val fold_pattern + : f:('a * 'r1 First.Pattern.t -> 'r1) + -> g:('b * ('r1, 'r2) Pattern.t -> 'r2) -> ('a, 'b) t + -> 'r2 + (** `rewrite_bottom_up` specializes `fold_pattern` so that the result type `'r1` is equal to the type of the nested fixed-point type i.e. `'r1 = 'a First.t` and the result type `'r2` is equal to the top-level @@ -144,35 +154,44 @@ module type S2 = sig fixed-point type `'a First.t` as its argument and `g` can be written with `('a,'b) t` as its argument. *) - - val unfold_pattern : - f:('r1 -> 'a * 'r1 First.Pattern.t) - -> g:('r2 -> 'b * ('r1, 'r2) Pattern.t) - -> 'r2 + val rewrite_bottom_up + : f:('a First.t -> 'a First.t) + -> g:(('a, 'b) t -> ('a, 'b) t) + -> ('a, 'b) t -> ('a, 'b) t + (** `unfold_pattern` takes a seed value of type `'r2` and uses the function `g` to generate a tuple of meta-data and a pattern with types `'r1` and `'r2`. The functions proceeds by recursively applying `g` to the contained values of type `'r2` and `f` to values of type `'r1` finishing when the pattern contains no values of type `'r1` or `'r2`. *) + val unfold_pattern + : f:('r1 -> 'a * 'r1 First.Pattern.t) + -> g:('r2 -> 'b * ('r1, 'r2) Pattern.t) + -> 'r2 + -> ('a, 'b) t - val rewrite_top_down : - f:('a First.t -> 'a First.t) + (** `rewrite_top_down` specializes `unfold_pattern` in a manner analogous to + how `rewrite_bottom_up` specializes `fold_pattern` *) + val rewrite_top_down + : f:('a First.t -> 'a First.t) -> g:(('a, 'b) t -> ('a, 'b) t) -> ('a, 'b) t -> ('a, 'b) t - (** `rewrite_top_down` specializes `unfold_pattern` in a manner analogous to - how `rewrite_bottom_up` specializes `fold_pattern` *) end module Make2 (First : S) (Pattern : Pattern.S2) : S2 with module First := First and module Pattern := Pattern = struct - type ('a, 'b) t = {pattern: ('a First.t, ('a, 'b) t) Pattern.t; meta: 'b} + type ('a, 'b) t = + { pattern : ('a First.t, ('a, 'b) t) Pattern.t + ; meta : 'b + } [@@deriving map, fold, compare, hash, sexp] - let rec pp f g ppf {pattern; meta} = + let rec pp f g ppf { pattern; meta } = Fmt.pf ppf {|%a%a|} g meta (Pattern.pp (First.pp f) (pp f g)) pattern + ;; include Foldable.Make2 (struct type nonrec ('a, 'b) t = ('a, 'b) t @@ -180,34 +199,31 @@ module Make2 (First : S) (Pattern : Pattern.S2) : let fold = fold end) - let rec fold_pattern ~f ~g {meta; pattern} = - let pattern' = - Pattern.map (First.fold_pattern ~f) (fold_pattern ~f ~g) pattern - in + let rec fold_pattern ~f ~g { meta; pattern } = + let pattern' = Pattern.map (First.fold_pattern ~f) (fold_pattern ~f ~g) pattern in g (meta, pattern') + ;; (** fold_pattern (Fn.compose f First.fix) (Fn.compose g fix) x *) let rec rewrite_bottom_up ~f ~g t = g { t with - pattern= - Pattern.map - (First.rewrite_bottom_up ~f) - (rewrite_bottom_up ~f ~g) t.pattern } + pattern = + Pattern.map (First.rewrite_bottom_up ~f) (rewrite_bottom_up ~f ~g) t.pattern + } + ;; let rec unfold_pattern ~f ~g x = let meta, pattern = g x in - let pattern' = - Pattern.map (First.unfold_pattern ~f) (unfold_pattern ~f ~g) pattern - in - {meta; pattern= pattern'} + let pattern' = Pattern.map (First.unfold_pattern ~f) (unfold_pattern ~f ~g) pattern in + { meta; pattern = pattern' } + ;; (**`unfold_pattern (Fn.compose First.unfix f) (Fn.compose unfix g) x`*) let rec rewrite_top_down ~f ~g x = let t = g x in { t with - pattern= - Pattern.map - (First.rewrite_top_down ~f) - (rewrite_top_down ~f ~g) t.pattern } + pattern = Pattern.map (First.rewrite_top_down ~f) (rewrite_top_down ~f ~g) t.pattern + } + ;; end diff --git a/src/common/Foldable.ml b/src/common/Foldable.ml index 1cf6166476..95c75454c5 100644 --- a/src/common/Foldable.ml +++ b/src/common/Foldable.ml @@ -33,22 +33,22 @@ module type S = sig (** A data structure which can be folded *) type 'a t - val fold_left : f:('b -> 'a -> 'b) -> init:'b -> 'a t -> 'b (** Left associative fold of a data structure; this is the same as the function derived from `[@@deriving fold]` but with labelled arguments *) + val fold_left : f:('b -> 'a -> 'b) -> init:'b -> 'a t -> 'b - val fold_right : f:('a -> 'b -> 'b) -> init:'b -> 'a t -> 'b (** Right associative fold of a data structure *) + val fold_right : f:('a -> 'b -> 'b) -> init:'b -> 'a t -> 'b - val any : pred:('a -> bool) -> ?init:bool -> 'a t -> bool (** Test whether any element of the data structure satisfies the supplied predicate. The optional argument `init` specifies the starting value and defaults to `false`. *) + val any : pred:('a -> bool) -> ?init:bool -> 'a t -> bool - val all : pred:('a -> bool) -> ?init:bool -> 'a t -> bool (** Test whether all elements of the the data structure satify the supplied predicate. The optional argument `init` specifies the starting value and defaults to `true`. *) + val all : pred:('a -> bool) -> ?init:bool -> 'a t -> bool end module Make (X : Basic) : S with type 'a t := 'a X.t = struct @@ -58,12 +58,10 @@ module Make (X : Basic) : S with type 'a t := 'a X.t = struct let fold_right ~f ~init x = let f' k x z = k @@ f x z in fold_left ~f:f' ~init:(fun x -> x) x init + ;; - let any ~pred ?(init = false) x = - fold_right ~f:(fun x accu -> accu || pred x) ~init x - - let all ~pred ?(init = true) x = - fold_right ~f:(fun x accu -> accu && pred x) ~init x + let any ~pred ?(init = false) x = fold_right ~f:(fun x accu -> accu || pred x) ~init x + let all ~pred ?(init = true) x = fold_right ~f:(fun x accu -> accu && pred x) ~init x end (* The `Basic` definition for type constructors with two type variables. *) @@ -77,21 +75,18 @@ type variable. module type S2 = sig type ('a, 'b) t - val fold_left : - f:('c -> 'a -> 'c) -> g:('c -> 'b -> 'c) -> init:'c -> ('a, 'b) t -> 'c - - val fold_right : - f:('a -> 'c -> 'c) -> g:('b -> 'c -> 'c) -> init:'c -> ('a, 'b) t -> 'c + val fold_left : f:('c -> 'a -> 'c) -> g:('c -> 'b -> 'c) -> init:'c -> ('a, 'b) t -> 'c + val fold_right : f:('a -> 'c -> 'c) -> g:('b -> 'c -> 'c) -> init:'c -> ('a, 'b) t -> 'c - val any : - pred_first:('a -> bool) + val any + : pred_first:('a -> bool) -> pred_second:('b -> bool) -> ?init:bool -> ('a, 'b) t -> bool - val all : - pred_first:('a -> bool) + val all + : pred_first:('a -> bool) -> pred_second:('b -> bool) -> ?init:bool -> ('a, 'b) t @@ -102,18 +97,24 @@ module Make2 (X : Basic2) : S2 with type ('a, 'b) t := ('a, 'b) X.t = struct let fold_left ~f ~g ~init x = X.fold f g init x let fold_right ~f ~g ~init x = - let f' k x z = k @@ f x z and g' k x z = k @@ g x z in + let f' k x z = k @@ f x z + and g' k x z = k @@ g x z in fold_left ~f:f' ~g:g' ~init:Fn.id x init + ;; let any ~pred_first ~pred_second ?(init = false) x = fold_right ~f:(fun x accu -> accu || pred_first x) ~g:(fun x accu -> accu || pred_second x) - ~init x + ~init + x + ;; let all ~pred_first ~pred_second ?(init = true) x = fold_right ~f:(fun x accu -> accu && pred_first x) ~g:(fun x accu -> accu && pred_second x) - ~init x + ~init + x + ;; end diff --git a/src/common/Gensym.ml b/src/common/Gensym.ml index ba09069679..f6b6ea4c77 100644 --- a/src/common/Gensym.ml +++ b/src/common/Gensym.ml @@ -1,11 +1,13 @@ let _counter = ref 0 let generate ?(prefix : string = "") () = - _counter := !_counter + 1 ; + _counter := !_counter + 1; Format.sprintf "%ssym%d__" prefix !_counter +;; let enter () = let old_counter = !_counter in - (generate (), fun () -> _counter := old_counter) + generate (), fun () -> _counter := old_counter +;; let reset_danger_use_cautiously () = _counter := 0 diff --git a/src/common/Helpers.ml b/src/common/Helpers.ml index fd8161267e..036667f8e5 100644 --- a/src/common/Helpers.ml +++ b/src/common/Helpers.ml @@ -1,8 +1,8 @@ open Core_kernel let option_or_else ~if_none x = Option.first_some x if_none -let on_snd f (x, y) = (x, f y) -let on_fst f (x, y) = (f x, y) +let on_snd f (x, y) = x, f y +let on_fst f (x, y) = f x, y let curry f x y = f (x, y) let uncurry f (x, y) = f x y let pp_builtin_syntax = Fmt.(string |> styled `Yellow) diff --git a/src/common/Label.ml b/src/common/Label.ml index 0d027ec75a..6109fc739a 100644 --- a/src/common/Label.ml +++ b/src/common/Label.ml @@ -7,11 +7,7 @@ module type S = sig include Pretty.S with type t := t include Comparator.S with type t := t - - include - Comparable.S - with type t := t - and type comparator_witness := comparator_witness + include Comparable.S with type t := t and type comparator_witness := comparator_witness val init : t val next : t -> t diff --git a/src/common/Pretty.ml b/src/common/Pretty.ml index 8eb6f927ce..0fbd19d320 100644 --- a/src/common/Pretty.ml +++ b/src/common/Pretty.ml @@ -16,8 +16,8 @@ end module type S2 = sig type ('a, 'b) t - val pp : - (Format.formatter -> 'a -> unit) + val pp + : (Format.formatter -> 'a -> unit) -> (Format.formatter -> 'b -> unit) -> Format.formatter -> ('a, 'b) t diff --git a/src/common/Specialized.ml b/src/common/Specialized.ml index 1a11fc2226..e21ff8e2ed 100644 --- a/src/common/Specialized.ml +++ b/src/common/Specialized.ml @@ -39,16 +39,11 @@ module type S = sig module Meta : Meta include Pretty.S with type t := t include Comparator.S with type t := t - - include - Comparable.S - with type t := t - and type comparator_witness := comparator_witness + include Comparable.S with type t := t and type comparator_witness := comparator_witness end module Make (X : Unspecialized) (Meta : Meta) : - S with type t = (Meta.t[@compare.ignore]) X.t and module Meta := Meta = -struct + S with type t = (Meta.t[@compare.ignore]) X.t and module Meta := Meta = struct module Basic = struct type t = (Meta.t[@compare.ignore]) X.t [@@deriving hash, sexp, compare] @@ -68,9 +63,8 @@ end module Make2 (X : Unspecialized2) (First : S) (Meta : Meta) : S - with type t = - ((First.Meta.t[@compare.ignore]), (Meta.t[@compare.ignore])) X.t - and module Meta := Meta = struct + with type t = ((First.Meta.t[@compare.ignore]), (Meta.t[@compare.ignore])) X.t + and module Meta := Meta = struct module Basic = struct type t = ((First.Meta.t[@compare.ignore]), (Meta.t[@compare.ignore])) X.t [@@deriving hash, sexp, compare] diff --git a/src/common/Validation.ml b/src/common/Validation.ml index 6f94a53af8..e8be72a740 100644 --- a/src/common/Validation.ml +++ b/src/common/Validation.ml @@ -25,10 +25,7 @@ module type S = sig val get_errors_opt : 'a t -> error list option val get_first_error_opt : 'a t -> error option val get_success_opt : 'a t -> 'a option - - val get_with : - 'a t -> with_ok:('a -> 'b) -> with_errors:(error list -> 'b) -> 'b - + val get_with : 'a t -> with_ok:('a -> 'b) -> with_errors:(error list -> 'b) -> 'b val to_result : 'a t -> ('a, error list) result module Validation_infix : Infix with type 'a t := 'a t @@ -116,52 +113,85 @@ end) : S with type error := X.t = struct type errors = NonEmpty of X.t * X.t list let append xs ys = - match (xs, ys) with - | NonEmpty (x, []), NonEmpty (y, []) -> NonEmpty (x, [y]) + match xs, ys with + | NonEmpty (x, []), NonEmpty (y, []) -> NonEmpty (x, [ y ]) | NonEmpty (x, xs), NonEmpty (y, ys) -> NonEmpty (x, xs @ (y :: ys)) + ;; type 'a t = ('a, errors) result - let map x ~f = match x with Ok x -> Ok (f x) | Error x -> Error x + let map x ~f = + match x with + | Ok x -> Ok (f x) + | Error x -> Error x + ;; + let pure x = Ok x let apply x ~f = - match (f, x) with + match f, x with | Ok f, Ok x -> Ok (f x) | Error e, Ok _ -> Error e | Ok _, Error e -> Error e | Error e1, Error e2 -> Error (append e1 e2) + ;; let apply_const a b = apply b ~f:(map a ~f:(fun _ x -> x)) - let bind x ~f = match x with Ok x -> f x | Error e -> Error e + + let bind x ~f = + match x with + | Ok x -> f x + | Error e -> Error e + ;; + let liftA2 f x y = apply y ~f:(apply x ~f:(pure f)) let liftA3 f x y z = apply z ~f:(apply y ~f:(apply x ~f:(pure f))) let consA next rest = liftA2 List.cons next rest let sequence ts = List.fold_right ~init:(pure []) ~f:consA ts - module Validation_infix = struct let ( >>= ) x f = bind x ~f end + module Validation_infix = struct + let ( >>= ) x f = bind x ~f + end + include Validation_infix let ok x = pure x let error x = Error (NonEmpty (x, [])) - let is_error = function Error _ -> true | _ -> false - let is_success = function Ok _ -> true | _ -> false + + let is_error = function + | Error _ -> true + | _ -> false + ;; + + let is_success = function + | Ok _ -> true + | _ -> false + ;; let get_errors_opt = function | Error (NonEmpty (x, xs)) -> Some (x :: xs) | _ -> None + ;; let get_first_error_opt = function | Error (NonEmpty (x, _)) -> Some x | _ -> None + ;; - let get_success_opt = function Ok x -> Some x | _ -> None + let get_success_opt = function + | Ok x -> Some x + | _ -> None + ;; let get_with x ~with_ok ~with_errors = match x with | Ok x -> with_ok x | Error (NonEmpty (x, xs)) -> with_errors @@ (x :: xs) + ;; let to_result x = - match x with Ok x -> Ok x | Error (NonEmpty (x, xs)) -> Error (x :: xs) + match x with + | Ok x -> Ok x + | Error (NonEmpty (x, xs)) -> Error (x :: xs) + ;; end diff --git a/src/common/Validation.mli b/src/common/Validation.mli index d158a3251c..516f7b112d 100644 --- a/src/common/Validation.mli +++ b/src/common/Validation.mli @@ -23,10 +23,7 @@ module type S = sig val get_errors_opt : 'a t -> error list option val get_first_error_opt : 'a t -> error option val get_success_opt : 'a t -> 'a option - - val get_with : - 'a t -> with_ok:('a -> 'b) -> with_errors:(error list -> 'b) -> 'b - + val get_with : 'a t -> with_ok:('a -> 'b) -> with_errors:(error list -> 'b) -> 'b val to_result : 'a t -> ('a, error list) result module Validation_infix : Infix with type 'a t := 'a t diff --git a/src/frontend/Ast.ml b/src/frontend/Ast.ml index 80f5e1fd34..84b078775f 100644 --- a/src/frontend/Ast.ml +++ b/src/frontend/Ast.ml @@ -5,7 +5,9 @@ open Middle (** Our type for identifiers, on which we record a location *) type identifier = - {name: string; id_loc: Location_span.t sexp_opaque [@compare.ignore]} + { name : string + ; id_loc : Location_span.t sexp_opaque [@compare.ignore] + } [@@deriving sexp, hash, compare] (** Indices for array access *) @@ -18,7 +20,10 @@ type 'e index = [@@deriving sexp, hash, compare, map] (** Front-end function kinds *) -type fun_kind = StanLib | UserDefined [@@deriving compare, sexp, hash] +type fun_kind = + | StanLib + | UserDefined +[@@deriving compare, sexp, hash] (** Expression shapes (used for both typed and untyped expressions, where we substitute untyped_expression or typed_expression for 'e *) @@ -41,11 +46,14 @@ type ('e, 'f) expression = | Indexed of 'e * 'e index list [@@deriving sexp, hash, compare, map] -type ('m, 'f) expr_with = {expr: (('m, 'f) expr_with, 'f) expression; emeta: 'm} +type ('m, 'f) expr_with = + { expr : (('m, 'f) expr_with, 'f) expression + ; emeta : 'm + } [@@deriving sexp, compare, map, hash] (** Untyped expressions, which have location_spans as meta-data *) -type located_meta = {loc: Location_span.t sexp_opaque [@compare.ignore]} +type located_meta = { loc : Location_span.t sexp_opaque [@compare.ignore] } [@@deriving sexp, compare, map, hash] type untyped_expression = (located_meta, unit) expr_with @@ -54,28 +62,32 @@ type untyped_expression = (located_meta, unit) expr_with (** Typed expressions also have meta-data after type checking: a location_span, as well as a type and an origin block (lub of the origin blocks of the identifiers in it) *) type typed_expr_meta = - { loc: Location_span.t sexp_opaque [@compare.ignore] - ; ad_level: UnsizedType.autodifftype - ; type_: UnsizedType.t } + { loc : Location_span.t sexp_opaque [@compare.ignore] + ; ad_level : UnsizedType.autodifftype + ; type_ : UnsizedType.t + } [@@deriving sexp, compare, map, hash] type typed_expression = (typed_expr_meta, fun_kind) expr_with [@@deriving sexp, compare, map, hash] -let mk_untyped_expression ~expr ~loc = {expr; emeta= {loc}} +let mk_untyped_expression ~expr ~loc = { expr; emeta = { loc } } let mk_typed_expression ~expr ~loc ~type_ ~ad_level = - {expr; emeta= {loc; type_; ad_level}} + { expr; emeta = { loc; type_; ad_level } } +;; let expr_loc_lub exprs = match List.map ~f:(fun e -> e.emeta.loc) exprs with | [] -> raise_s [%message "Can't find location lub for empty list"] - | [hd] -> hd + | [ hd ] -> hd | x1 :: tl -> List.fold ~init:x1 ~f:Location_span.merge tl +;; (** Least upper bound of expression autodiff types *) let expr_ad_lub exprs = exprs |> List.map ~f:(fun x -> x.emeta.ad_level) |> UnsizedType.lub_ad_type +;; (** Assignment operators *) type assignmentoperator = @@ -94,7 +106,9 @@ type 'e truncation = [@@deriving sexp, hash, compare, map] (** Things that can be printed *) -type 'e printable = PString of string | PExpr of 'e +type 'e printable = + | PString of string + | PExpr of 'e [@@deriving sexp, compare, map, hash] type ('l, 'e) lvalue = @@ -102,7 +116,10 @@ type ('l, 'e) lvalue = | LIndexed of 'l * 'e index list [@@deriving sexp, hash, compare, map] -type ('e, 'm) lval_with = {lval: (('e, 'm) lval_with, 'e) lvalue; lmeta: 'm} +type ('e, 'm) lval_with = + { lval : (('e, 'm) lval_with, 'e) lvalue + ; lmeta : 'm + } [@@deriving sexp, hash, compare, map] type untyped_lval = (untyped_expression, located_meta) lval_with @@ -116,18 +133,20 @@ type typed_lval = (typed_expression, typed_expr_meta) lval_with typed_statement to get typed_statement *) type ('e, 's, 'l, 'f) statement = | Assignment of - { assign_lhs: 'l - ; assign_op: assignmentoperator - ; assign_rhs: 'e } + { assign_lhs : 'l + ; assign_op : assignmentoperator + ; assign_rhs : 'e + } | NRFunApp of 'f * identifier * 'e list | TargetPE of 'e (* IncrementLogProb is deprecated *) | IncrementLogProb of 'e | Tilde of - { arg: 'e - ; distribution: identifier - ; args: 'e list - ; truncation: 'e truncation } + { arg : 'e + ; distribution : identifier + ; args : 'e list + ; truncation : 'e truncation + } | Break | Continue | Return of 'e @@ -138,25 +157,27 @@ type ('e, 's, 'l, 'f) statement = | IfThenElse of 'e * 's * 's option | While of 'e * 's | For of - { loop_variable: identifier - ; lower_bound: 'e - ; upper_bound: 'e - ; loop_body: 's } + { loop_variable : identifier + ; lower_bound : 'e + ; upper_bound : 'e + ; loop_body : 's + } | ForEach of identifier * 'e * 's | Block of 's list | VarDecl of - { decl_type: 'e Middle.Type.t - ; transformation: 'e Middle.Program.transformation - ; identifier: identifier - ; initial_value: 'e option - ; is_global: bool } + { decl_type : 'e Middle.Type.t + ; transformation : 'e Middle.Program.transformation + ; identifier : identifier + ; initial_value : 'e option + ; is_global : bool + } | FunDef of - { returntype: Middle.UnsizedType.returntype - ; funname: identifier - ; arguments: - (Middle.UnsizedType.autodifftype * Middle.UnsizedType.t * identifier) - list - ; body: 's } + { returntype : Middle.UnsizedType.returntype + ; funname : identifier + ; arguments : + (Middle.UnsizedType.autodifftype * Middle.UnsizedType.t * identifier) list + ; body : 's + } [@@deriving sexp, hash, compare, map] (** Statement return types which we will decorate statements with during type @@ -175,7 +196,9 @@ type statement_returntype = [@@deriving sexp, hash, compare] type ('e, 'm, 'l, 'f) statement_with = - {stmt: ('e, ('e, 'm, 'l, 'f) statement_with, 'l, 'f) statement; smeta: 'm} + { stmt : ('e, ('e, 'm, 'l, 'f) statement_with, 'l, 'f) statement + ; smeta : 'm + } [@@deriving sexp, compare, map, hash] (** Untyped statements, which have location_spans as meta-data *) @@ -183,41 +206,37 @@ type untyped_statement = (untyped_expression, located_meta, untyped_lval, unit) statement_with [@@deriving sexp, compare, map, hash] -let mk_untyped_statement ~stmt ~loc : untyped_statement = {stmt; smeta= {loc}} +let mk_untyped_statement ~stmt ~loc : untyped_statement = { stmt; smeta = { loc } } type stmt_typed_located_meta = - { loc: Middle.Location_span.t sexp_opaque [@compare.ignore] - ; return_type: statement_returntype } + { loc : Middle.Location_span.t sexp_opaque [@compare.ignore] + ; return_type : statement_returntype + } [@@deriving sexp, compare, map, hash] (** Typed statements also have meta-data after type checking: a location_span, as well as a statement returntype to check that function bodies have the right return type*) type typed_statement = - ( typed_expression - , stmt_typed_located_meta - , typed_lval - , fun_kind ) - statement_with + (typed_expression, stmt_typed_located_meta, typed_lval, fun_kind) statement_with [@@deriving sexp, compare, map, hash] -let mk_typed_statement ~stmt ~loc ~return_type = - {stmt; smeta= {loc; return_type}} +let mk_typed_statement ~stmt ~loc ~return_type = { stmt; smeta = { loc; return_type } } (** Program shapes, where we obtain types of programs if we substitute typed or untyped statements for 's *) type 's program = - { functionblock: 's list option - ; datablock: 's list option - ; transformeddatablock: 's list option - ; parametersblock: 's list option - ; transformedparametersblock: 's list option - ; modelblock: 's list option - ; generatedquantitiesblock: 's list option } + { functionblock : 's list option + ; datablock : 's list option + ; transformeddatablock : 's list option + ; parametersblock : 's list option + ; transformedparametersblock : 's list option + ; modelblock : 's list option + ; generatedquantitiesblock : 's list option + } [@@deriving sexp, hash, compare, map] (** Untyped programs (before type checking) *) -type untyped_program = untyped_statement program -[@@deriving sexp, compare, map] +type untyped_program = untyped_statement program [@@deriving sexp, compare, map] (** Typed programs (after type checking) *) type typed_program = typed_statement program [@@deriving sexp, compare, map] @@ -225,46 +244,63 @@ type typed_program = typed_statement program [@@deriving sexp, compare, map] (*========================== Helper functions ===============================*) (** 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} } - -let rec untyped_lvalue_of_typed_lvalue ({lval; lmeta} : typed_lval) : - untyped_lval = - { lval= - map_lvalue untyped_lvalue_of_typed_lvalue - untyped_expression_of_typed_expression lval - ; lmeta= {loc= lmeta.loc} } +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 } + } +;; + +let rec untyped_lvalue_of_typed_lvalue ({ lval; lmeta } : typed_lval) : untyped_lval = + { lval = + map_lvalue + untyped_lvalue_of_typed_lvalue + untyped_expression_of_typed_expression + lval + ; lmeta = { loc = lmeta.loc } + } +;; (** Forgetful function from typed to untyped statements *) -let rec untyped_statement_of_typed_statement {stmt; smeta} = - { stmt= - map_statement untyped_expression_of_typed_expression - untyped_statement_of_typed_statement untyped_lvalue_of_typed_lvalue +let rec untyped_statement_of_typed_statement { stmt; smeta } = + { stmt = + map_statement + untyped_expression_of_typed_expression + untyped_statement_of_typed_statement + untyped_lvalue_of_typed_lvalue (fun _ -> ()) stmt - ; smeta= {loc= smeta.loc} } + ; smeta = { loc = smeta.loc } + } +;; (** Forgetful function from typed to untyped programs *) let untyped_program_of_typed_program : typed_program -> untyped_program = map_program untyped_statement_of_typed_statement +;; -let rec expr_of_lvalue {lval; lmeta} = - { expr= - ( match lval with +let rec expr_of_lvalue { lval; lmeta } = + { expr = + (match lval with | LVariable s -> Variable s - | LIndexed (l, i) -> Indexed (expr_of_lvalue l, i) ) - ; emeta= lmeta } - -let rec lvalue_of_expr {expr; emeta} = - { lval= - ( match expr with + | LIndexed (l, i) -> Indexed (expr_of_lvalue l, i)) + ; emeta = lmeta + } +;; + +let rec lvalue_of_expr { expr; emeta } = + { lval = + (match expr with | Variable s -> LVariable s | Indexed (l, i) -> LIndexed (lvalue_of_expr l, i) - | _ -> failwith "Trying to convert illegal expression to lval." ) - ; lmeta= emeta } - -let rec id_of_lvalue {lval; _} = - match lval with LVariable s -> s | LIndexed (l, _) -> id_of_lvalue l + | _ -> failwith "Trying to convert illegal expression to lval.") + ; lmeta = emeta + } +;; + +let rec id_of_lvalue { lval; _ } = + match lval with + | LVariable s -> s + | LIndexed (l, _) -> id_of_lvalue l +;; diff --git a/src/frontend/Ast_to_Mir.ml b/src/frontend/Ast_to_Mir.ml index db858ba754..101f4dab55 100644 --- a/src/frontend/Ast_to_Mir.ml +++ b/src/frontend/Ast_to_Mir.ml @@ -4,100 +4,92 @@ open Middle (* XXX fix exn *) let unwrap_return_exn = function | Some (UnsizedType.ReturnType ut) -> ut - | x -> - raise_s - [%message - "Unexpected return type " (x : UnsizedType.returntype option)] + | x -> raise_s [%message "Unexpected return type " (x : UnsizedType.returntype option)] +;; let trans_fn_kind = function | Ast.StanLib -> Fun_kind.StanLib | UserDefined -> UserDefined +;; let without_underscores = String.filter ~f:(( <> ) '_') let drop_leading_zeros s = match String.lfindi ~f:(fun _ c -> c <> '0') s with - | Some p when p > 0 -> ( - match s.[p] with + | Some p when p > 0 -> + (match s.[p] with | 'e' | '.' -> String.drop_prefix s (p - 1) - | _ -> String.drop_prefix s p ) + | _ -> String.drop_prefix s p) | Some _ -> s | None -> "0" +;; let format_number s = s |> without_underscores |> drop_leading_zeros let%expect_test "format_number0" = - format_number "0_000." |> print_endline ; + format_number "0_000." |> print_endline; [%expect "0."] +;; let%expect_test "format_number1" = - format_number ".123_456" |> print_endline ; + format_number ".123_456" |> print_endline; [%expect ".123456"] +;; let rec op_to_funapp op args = - let argtypes = - List.map ~f:(fun x -> (x.Ast.emeta.Ast.ad_level, x.emeta.type_)) args - in + let argtypes = List.map ~f:(fun x -> x.Ast.emeta.Ast.ad_level, x.emeta.type_) args in let type_ = - Stan_math_signatures.operator_stan_math_return_type op argtypes - |> unwrap_return_exn + Stan_math_signatures.operator_stan_math_return_type op argtypes |> unwrap_return_exn and loc = Ast.expr_loc_lub args and adlevel = Ast.expr_ad_lub args in Expr. - { Fixed.pattern= FunApp (StanLib, Operator.to_string op, trans_exprs args) - ; meta= Expr.Typed.Meta.create ~type_ ~adlevel ~loc () } + { Fixed.pattern = FunApp (StanLib, Operator.to_string op, trans_exprs args) + ; meta = Expr.Typed.Meta.create ~type_ ~adlevel ~loc () + } -and trans_expr {Ast.expr; Ast.emeta} = +and trans_expr { Ast.expr; Ast.emeta } = let ewrap pattern = Expr. { Fixed.pattern - ; meta= + ; meta = Typed.Meta. - {type_= emeta.Ast.type_; adlevel= emeta.ad_level; loc= emeta.loc} + { type_ = emeta.Ast.type_; adlevel = emeta.ad_level; loc = emeta.loc } } in match expr with | Ast.Paren x -> trans_expr x | BinOp (lhs, And, rhs) -> EAnd (trans_expr lhs, trans_expr rhs) |> ewrap | BinOp (lhs, Or, rhs) -> EOr (trans_expr lhs, trans_expr rhs) |> ewrap - | BinOp (lhs, op, rhs) -> op_to_funapp op [lhs; rhs] - | PrefixOp (op, e) | Ast.PostfixOp (e, op) -> op_to_funapp op [e] + | BinOp (lhs, op, rhs) -> op_to_funapp op [ lhs; rhs ] + | PrefixOp (op, e) | Ast.PostfixOp (e, op) -> op_to_funapp op [ e ] | Ast.TernaryIf (cond, ifb, elseb) -> - Expr.Fixed.Pattern.TernaryIf - (trans_expr cond, trans_expr ifb, trans_expr elseb) - |> ewrap - | Variable {name; _} -> Var name |> ewrap + Expr.Fixed.Pattern.TernaryIf (trans_expr cond, trans_expr ifb, trans_expr elseb) + |> ewrap + | Variable { name; _ } -> Var name |> ewrap | IntNumeral x -> Lit (Int, format_number x) |> ewrap | RealNumeral x -> Lit (Real, format_number x) |> ewrap - | FunApp (fn_kind, {name; _}, args) | CondDistApp (fn_kind, {name; _}, args) - -> - FunApp (trans_fn_kind fn_kind, name, trans_exprs args) |> ewrap + | FunApp (fn_kind, { name; _ }, args) | CondDistApp (fn_kind, { name; _ }, args) -> + FunApp (trans_fn_kind fn_kind, name, trans_exprs args) |> ewrap | GetLP | GetTarget -> FunApp (StanLib, "target", []) |> ewrap | ArrayExpr eles -> - FunApp - (CompilerInternal, Internal_fun.to_string FnMakeArray, trans_exprs eles) - |> ewrap + FunApp (CompilerInternal, Internal_fun.to_string FnMakeArray, trans_exprs eles) + |> ewrap | RowVectorExpr eles -> - FunApp - ( CompilerInternal - , Internal_fun.to_string FnMakeRowVec - , trans_exprs eles ) - |> ewrap + FunApp (CompilerInternal, Internal_fun.to_string FnMakeRowVec, trans_exprs eles) + |> ewrap | Indexed (lhs, indices) -> - Indexed (trans_expr lhs, List.map ~f:trans_idx indices) |> ewrap + Indexed (trans_expr lhs, List.map ~f:trans_idx indices) |> ewrap and trans_idx = function | Ast.All -> All | Ast.Upfrom e -> Upfrom (trans_expr e) | Ast.Downfrom e -> Between (Expr.Helpers.loop_bottom, trans_expr e) | Ast.Between (lb, ub) -> Between (trans_expr lb, trans_expr ub) - | Ast.Single e -> ( - match e.emeta.type_ with + | Ast.Single e -> + (match e.emeta.type_ with | UInt -> Single (trans_expr e) | UArray _ -> MultiIndex (trans_expr e) - | _ -> - raise_s - [%message "Expecting int or array" (e.emeta.type_ : UnsizedType.t)] ) + | _ -> raise_s [%message "Expecting int or array" (e.emeta.type_ : UnsizedType.t)]) and trans_exprs exprs = List.map ~f:trans_expr exprs @@ -106,87 +98,102 @@ let trans_possiblysizedtype pst = Type.map trans_expr pst let neg_inf = Expr. - { Fixed.pattern= FunApp (StanLib, Internal_fun.to_string FnNegInf, []) - ; meta= - Typed.Meta.{type_= UReal; loc= Location_span.empty; adlevel= DataOnly} + { Fixed.pattern = FunApp (StanLib, Internal_fun.to_string FnNegInf, []) + ; meta = Typed.Meta.{ type_ = UReal; loc = Location_span.empty; adlevel = DataOnly } } +;; -let trans_arg (adtype, ut, ident) = (adtype, ident.Ast.name, ut) +let trans_arg (adtype, ut, ident) = adtype, ident.Ast.name, ut let truncate_dist ud_dists (id : Ast.identifier) ast_obs ast_args t = - let cdf_suffices = ["_lcdf"; "_cdf_log"] in - let ccdf_suffices = ["_lccdf"; "_ccdf_log"] in + let cdf_suffices = [ "_lcdf"; "_cdf_log" ] in + let ccdf_suffices = [ "_lccdf"; "_ccdf_log" ] in let find_function_info sfx = - let possible_names = - List.map ~f:(( ^ ) id.name) sfx |> String.Set.of_list - in + let possible_names = List.map ~f:(( ^ ) id.name) sfx |> String.Set.of_list in match List.find ~f:(fun (n, _) -> Set.mem possible_names n) ud_dists with - | Some (name, tp) -> (Ast.UserDefined, name, tp) + | Some (name, tp) -> Ast.UserDefined, name, tp | None -> - ( Ast.StanLib - , Set.to_list possible_names |> List.hd_exn - , if Stan_math_signatures.is_stan_math_function_name (id.name ^ "_lpmf") - then UnsizedType.UInt - else UnsizedType.UReal (* close enough *) ) + ( Ast.StanLib + , Set.to_list possible_names |> List.hd_exn + , if Stan_math_signatures.is_stan_math_function_name (id.name ^ "_lpmf") + then UnsizedType.UInt + else UnsizedType.UReal (* close enough *) ) in let trunc cond_op (x : Ast.typed_expression) y = let smeta = x.Ast.emeta.loc in - { Stmt.Fixed.meta= smeta - ; pattern= + { Stmt.Fixed.meta = smeta + ; pattern = IfElse - ( op_to_funapp cond_op [ast_obs; x] - , {Stmt.Fixed.meta= smeta; pattern= TargetPE neg_inf} - , Some y ) } + ( op_to_funapp cond_op [ ast_obs; x ] + , { Stmt.Fixed.meta = smeta; pattern = TargetPE neg_inf } + , Some y ) + } in let targetme loc e = - {Stmt.Fixed.meta= loc; pattern= TargetPE (op_to_funapp Operator.PMinus [e])} + { Stmt.Fixed.meta = loc; pattern = TargetPE (op_to_funapp Operator.PMinus [ e ]) } in let funapp meta kind name args = - { Ast.emeta= meta - ; expr= Ast.FunApp (kind, {name; id_loc= Location_span.empty}, args) } + { Ast.emeta = meta + ; expr = Ast.FunApp (kind, { name; id_loc = Location_span.empty }, args) + } in let inclusive_bound tp (lb : Ast.typed_expression) = let emeta = lb.emeta in - if UnsizedType.is_int_type tp then + if UnsizedType.is_int_type tp + then Ast. - { emeta - ; expr= BinOp (lb, Operator.Minus, {emeta; expr= Ast.IntNumeral "1"}) - } + { emeta; expr = BinOp (lb, Operator.Minus, { emeta; expr = Ast.IntNumeral "1" }) } else lb in match t with | Ast.NoTruncate -> [] | TruncateUpFrom lb -> - let fk, fn, tp = find_function_info ccdf_suffices in - [ trunc Less lb - (targetme lb.emeta.loc - (funapp lb.emeta fk fn (inclusive_bound tp lb :: ast_args))) ] + let fk, fn, tp = find_function_info ccdf_suffices in + [ trunc + Less + lb + (targetme + lb.emeta.loc + (funapp lb.emeta fk fn (inclusive_bound tp lb :: ast_args))) + ] | TruncateDownFrom ub -> - let fk, fn, _ = find_function_info cdf_suffices in - [ trunc Greater ub - (targetme ub.emeta.loc (funapp ub.emeta fk fn (ub :: ast_args))) ] + let fk, fn, _ = find_function_info cdf_suffices in + [ trunc Greater ub (targetme ub.emeta.loc (funapp ub.emeta fk fn (ub :: ast_args))) ] | TruncateBetween (lb, ub) -> - let fk, fn, tp = find_function_info cdf_suffices in - [ trunc Less lb - (trunc Greater ub - (targetme ub.emeta.loc - (funapp ub.emeta Ast.StanLib "log_diff_exp" - [ funapp ub.emeta fk fn (ub :: ast_args) - ; funapp ub.emeta fk fn (inclusive_bound tp lb :: ast_args) - ]))) ] + let fk, fn, tp = find_function_info cdf_suffices in + [ trunc + Less + lb + (trunc + Greater + ub + (targetme + ub.emeta.loc + (funapp + ub.emeta + Ast.StanLib + "log_diff_exp" + [ funapp ub.emeta fk fn (ub :: ast_args) + ; funapp ub.emeta fk fn (inclusive_bound tp lb :: ast_args) + ]))) + ] +;; let unquote s = - if s.[0] = '"' && s.[String.length s - 1] = '"' then - String.drop_suffix (String.drop_prefix s 1) 1 + if s.[0] = '"' && s.[String.length s - 1] = '"' + then String.drop_suffix (String.drop_prefix s 1) 1 else s +;; (* hack(sean): strings aren't real XXX add UString to MIR and maybe AST. *) let mkstring loc s = Expr. - { Fixed.pattern= Lit (Str, s) - ; meta= Typed.Meta.create ~type_:UReal ~loc ~adlevel:DataOnly () } + { Fixed.pattern = Lit (Str, s) + ; meta = Typed.Meta.create ~type_:UReal ~loc ~adlevel:DataOnly () + } +;; let trans_printables mloc (ps : Ast.typed_expression Ast.printable list) = List.map @@ -194,20 +201,28 @@ let trans_printables mloc (ps : Ast.typed_expression Ast.printable list) = | Ast.PString s -> mkstring mloc (unquote s) | Ast.PExpr e -> trans_expr e) ps +;; (* These types signal the context for a declaration during statement translation. They are only interpreted by trans_decl.*) -type constrainaction = Check | Constrain | Unconstrain [@@deriving sexp] +type constrainaction = + | Check + | Constrain + | Unconstrain +[@@deriving sexp] let constrainaction_fname c = Internal_fun.to_string - ( match c with + (match c with | Check -> FnCheck | Constrain -> FnConstrain - | Unconstrain -> FnUnconstrain ) + | Unconstrain -> FnUnconstrain) +;; type decl_context = - {dconstrain: constrainaction option; dadlevel: UnsizedType.autodifftype} + { dconstrain : constrainaction option + ; dadlevel : UnsizedType.autodifftype + } let check_constraint_to_string t (c : constrainaction) = match t with @@ -219,46 +234,63 @@ let check_constraint_to_string t (c : constrainaction) = | CholeskyCov -> "cholesky_factor" | Correlation -> "corr_matrix" | Covariance -> "cov_matrix" - | Lower _ -> ( - match c with + | 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" ) + | 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 constraint_forl = function - | Program.Identity | Offset _ | Multiplier _ | OffsetMultiplier _ | Lower _ - |Upper _ | LowerUpper _ -> - Stmt.Helpers.for_scalar - | Ordered | PositiveOrdered | Simplex | UnitVector | CholeskyCorr - |CholeskyCov | Correlation | Covariance -> - Stmt.Helpers.for_eigen + | Program.Identity + | Offset _ + | Multiplier _ + | OffsetMultiplier _ + | Lower _ + | Upper _ + | LowerUpper _ -> Stmt.Helpers.for_scalar + | Ordered + | PositiveOrdered + | Simplex + | UnitVector + | CholeskyCorr + | CholeskyCov + | Correlation + | Covariance -> Stmt.Helpers.for_eigen +;; let same_shape decl_id decl_var id var meta = - if UnsizedType.is_scalar_type (Expr.Typed.type_of var) then [] + if UnsizedType.is_scalar_type (Expr.Typed.type_of var) + then [] else [ Stmt. - { Fixed.pattern= + { Fixed.pattern = NRFunApp ( StanLib , "check_matching_dims" - , Expr.Helpers. - [str "constraint"; str decl_id; decl_var; str id; var] ) - ; meta } ] + , Expr.Helpers.[ str "constraint"; str decl_id; decl_var; str id; var ] ) + ; meta + } + ] +;; let check_transform_shape decl_id decl_var meta = function | Program.Offset e -> same_shape decl_id decl_var "offset" e meta @@ -266,48 +298,70 @@ let check_transform_shape decl_id decl_var meta = function | Lower e -> same_shape decl_id decl_var "lower" e meta | Upper e -> same_shape decl_id decl_var "upper" e meta | OffsetMultiplier (e1, e2) -> - same_shape decl_id decl_var "offset" e1 meta - @ same_shape decl_id decl_var "multiplier" e2 meta + same_shape decl_id decl_var "offset" e1 meta + @ same_shape decl_id decl_var "multiplier" e2 meta | LowerUpper (e1, e2) -> - same_shape decl_id decl_var "lower" e1 meta - @ same_shape decl_id decl_var "upper" e2 meta - | Covariance | Correlation | CholeskyCov | CholeskyCorr | Ordered - |PositiveOrdered | Simplex | UnitVector | Identity -> - [] + same_shape decl_id decl_var "lower" e1 meta + @ same_shape decl_id decl_var "upper" e2 meta + | Covariance + | Correlation + | CholeskyCov + | CholeskyCorr + | Ordered + | PositiveOrdered + | Simplex + | UnitVector + | Identity -> [] +;; let copy_indices indexed (var : Expr.Typed.t) = - if UnsizedType.is_scalar_type var.meta.type_ then var - else + if UnsizedType.is_scalar_type var.meta.type_ + then var + else ( match Expr.Helpers.collect_indices indexed with | [] -> var | indices -> - Expr.Fixed. - { pattern= Indexed (var, indices) - ; meta= - { var.meta with - type_= - Expr.Helpers.infer_type_of_indexed var.meta.type_ indices } - } + Expr.Fixed. + { pattern = Indexed (var, indices) + ; meta = + { var.meta with + type_ = Expr.Helpers.infer_type_of_indexed var.meta.type_ indices + } + }) +;; let extract_transform_args var = function - | Program.Lower a | Upper a -> [copy_indices var a] - | Offset a -> - [copy_indices var a; {a with Expr.Fixed.pattern= Lit (Int, "1")}] - | Multiplier a -> [{a with pattern= Lit (Int, "0")}; copy_indices var a] + | Program.Lower a | Upper a -> [ copy_indices var a ] + | Offset a -> [ copy_indices var a; { a with Expr.Fixed.pattern = Lit (Int, "1") } ] + | Multiplier a -> [ { a with pattern = Lit (Int, "0") }; copy_indices var a ] | LowerUpper (a1, a2) | OffsetMultiplier (a1, a2) -> - [copy_indices var a1; copy_indices var a2] - | Covariance | Correlation | CholeskyCov | CholeskyCorr | Ordered - |PositiveOrdered | Simplex | UnitVector | Identity -> - [] + [ copy_indices var a1; copy_indices var a2 ] + | Covariance + | Correlation + | CholeskyCov + | CholeskyCorr + | Ordered + | PositiveOrdered + | Simplex + | UnitVector + | Identity -> [] +;; let extra_constraint_args st = function - | Program.Lower _ | Upper _ | Offset _ | Multiplier _ | LowerUpper _ - |OffsetMultiplier _ | Ordered | PositiveOrdered | Simplex | UnitVector - |Identity -> - [] - | Covariance | Correlation | CholeskyCorr -> - [List.hd_exn (SizedType.dims_of st)] + | Program.Lower _ + | Upper _ + | Offset _ + | Multiplier _ + | LowerUpper _ + | OffsetMultiplier _ + | Ordered + | PositiveOrdered + | Simplex + | UnitVector + | Identity -> [] + | Covariance | Correlation | CholeskyCorr -> [ List.hd_exn (SizedType.dims_of st) ] | CholeskyCov -> SizedType.dims_of st +;; let param_size transform sizedtype = let rec shrink_eigen f st = @@ -315,53 +369,44 @@ let param_size transform sizedtype = | SizedType.SArray (t, d) -> SizedType.SArray (shrink_eigen f t, d) | SVector d | SMatrix (d, _) -> SVector (f d) | SInt | SReal | SRowVector _ -> - raise_s - [%message - "Expecting SVector or SMatrix, got " (st : Expr.Typed.t SizedType.t)] + raise_s + [%message "Expecting SVector or SMatrix, got " (st : Expr.Typed.t SizedType.t)] in let rec shrink_eigen_mat f st = match st with | SizedType.SArray (t, d) -> SizedType.SArray (shrink_eigen_mat f t, d) | SMatrix (d1, d2) -> SVector (f d1 d2) | SInt | SReal | SRowVector _ | SVector _ -> - raise_s - [%message "Expecting SMatrix, got " (st : Expr.Typed.t SizedType.t)] + raise_s [%message "Expecting SMatrix, got " (st : Expr.Typed.t SizedType.t)] in let k_choose_2 k = Expr.Helpers.(binop (binop k Times (binop k Minus (int 1))) Divide (int 2)) in match transform with | Program.Identity | Lower _ | Upper _ - |LowerUpper (_, _) - |Offset _ | Multiplier _ - |OffsetMultiplier (_, _) - |Ordered | PositiveOrdered | UnitVector -> - sizedtype - | Simplex -> - shrink_eigen (fun d -> Expr.Helpers.(binop d Minus (int 1))) sizedtype + | LowerUpper (_, _) + | Offset _ | Multiplier _ + | OffsetMultiplier (_, _) + | Ordered | PositiveOrdered | UnitVector -> sizedtype + | Simplex -> shrink_eigen (fun d -> Expr.Helpers.(binop d Minus (int 1))) sizedtype | CholeskyCorr | Correlation -> shrink_eigen k_choose_2 sizedtype | CholeskyCov -> - (* (N * (N + 1)) / 2 + (M - N) * N *) - shrink_eigen_mat - (fun m n -> - Expr.Helpers.( - binop - (binop (k_choose_2 n) Plus n) - Plus - (binop (binop m Minus n) Times n)) ) - sizedtype + (* (N * (N + 1)) / 2 + (M - N) * N *) + shrink_eigen_mat + (fun m n -> + Expr.Helpers.( + binop (binop (k_choose_2 n) Plus n) Plus (binop (binop m Minus n) Times n))) + sizedtype | Covariance -> - shrink_eigen - (fun k -> Expr.Helpers.(binop k Plus (k_choose_2 k))) - sizedtype + shrink_eigen (fun k -> Expr.Helpers.(binop k Plus (k_choose_2 k))) sizedtype +;; let remove_possibly_exn pst action loc = match pst with | Type.Sized st -> st | Unsized _ -> - raise_s - [%message - "Error extracting sizedtype" ~action ~loc:(loc : Location_span.t)] + raise_s [%message "Error extracting sizedtype" ~action ~loc:(loc : Location_span.t)] +;; let constrain_decl decl_type dconstrain t decl_id decl_var smeta = let st = remove_possibly_exn decl_type "constrain" smeta in @@ -369,40 +414,45 @@ let constrain_decl decl_type dconstrain t decl_id decl_var smeta = match Option.map ~f:(constrain_constraint_to_string t) dconstrain with | None | Some "" -> [] | Some constraint_str -> - let dc = Option.value_exn dconstrain in - let fname = constrainaction_fname dc in - let extra_args = - match dconstrain with - | Some Constrain -> extra_constraint_args st t - | _ -> [] - in - let args var = - (var :: mkstring constraint_str :: extract_transform_args var t) - @ extra_args - in - let constrainvar var = - { var with - Expr.Fixed.pattern= FunApp (CompilerInternal, fname, args var) } - in - let unconstrained_decls, decl_id, ut = - let ut = SizedType.to_unsized (param_size t st) in - match dconstrain with - | Some Unconstrain when SizedType.to_unsized st <> ut -> - ( [ Stmt.Fixed. - { pattern= - Decl - { decl_adtype= DataOnly - ; decl_id= decl_id ^ "_free__" - ; decl_type= Sized (param_size t st) } - ; meta= smeta } ] - , decl_id ^ "_free__" - , ut ) - | _ -> ([], decl_id, SizedType.to_unsized st) - in - unconstrained_decls - @ [ (constraint_forl t) st - (Stmt.Helpers.assign_indexed ut decl_id smeta constrainvar) - decl_var smeta ] + let dc = Option.value_exn dconstrain in + let fname = constrainaction_fname dc in + let extra_args = + match dconstrain with + | Some Constrain -> extra_constraint_args st t + | _ -> [] + in + let args var = + (var :: mkstring constraint_str :: extract_transform_args var t) @ extra_args + in + let constrainvar var = + { var with Expr.Fixed.pattern = FunApp (CompilerInternal, fname, args var) } + in + let unconstrained_decls, decl_id, ut = + let ut = SizedType.to_unsized (param_size t st) in + match dconstrain with + | Some Unconstrain when SizedType.to_unsized st <> ut -> + ( [ Stmt.Fixed. + { pattern = + Decl + { decl_adtype = DataOnly + ; decl_id = decl_id ^ "_free__" + ; decl_type = Sized (param_size t st) + } + ; meta = smeta + } + ] + , decl_id ^ "_free__" + , ut ) + | _ -> [], decl_id, SizedType.to_unsized st + in + unconstrained_decls + @ [ (constraint_forl t) + st + (Stmt.Helpers.assign_indexed ut decl_id smeta constrainvar) + decl_var + smeta + ] +;; let rec check_decl var decl_type' decl_id decl_trans smeta adlevel = let decl_type = remove_possibly_exn decl_type' "check" smeta in @@ -412,62 +462,63 @@ let rec check_decl var decl_type' decl_id decl_trans smeta adlevel = let args = extract_transform_args id decl_trans in Stmt.Helpers.internal_nrfunapp FnCheck (fn :: id_str :: id :: args) smeta in - [(constraint_forl decl_trans) decl_type check_id var smeta] + [ (constraint_forl decl_trans) decl_type check_id var smeta ] in match decl_trans with | Identity | Offset _ | Multiplier _ | OffsetMultiplier (_, _) -> [] | LowerUpper (lb, ub) -> - check_decl var decl_type' decl_id (Lower lb) smeta adlevel - @ check_decl var decl_type' decl_id (Upper ub) smeta adlevel + check_decl var decl_type' decl_id (Lower lb) smeta adlevel + @ check_decl var decl_type' decl_id (Upper ub) smeta adlevel | _ -> chk (mkstring smeta (check_constraint_to_string decl_trans Check)) var +;; -let trans_decl {dconstrain; dadlevel} smeta decl_type transform identifier - initial_value = +let trans_decl { dconstrain; dadlevel } smeta decl_type transform identifier initial_value + = let decl_id = identifier.Ast.name in let rhs = Option.map ~f:trans_expr initial_value in let dt = trans_possiblysizedtype decl_type in let decl_adtype = dadlevel in let decl_var = Expr. - { Fixed.pattern= Var decl_id - ; meta= - Typed.Meta.create ~adlevel:dadlevel ~loc:smeta + { Fixed.pattern = Var decl_id + ; meta = + Typed.Meta.create + ~adlevel:dadlevel + ~loc:smeta ~type_:(Type.to_unsized decl_type) - () } + () + } in let decl = - Stmt. - {Fixed.pattern= Decl {decl_adtype; decl_id; decl_type= dt}; meta= smeta} + Stmt.{ Fixed.pattern = Decl { decl_adtype; decl_id; decl_type = dt }; meta = smeta } in let rhs_assignment = Option.map ~f:(fun e -> - Stmt.Fixed. - {pattern= Assignment ((decl_id, e.meta.type_, []), e); meta= smeta} - ) + Stmt.Fixed.{ pattern = Assignment ((decl_id, e.meta.type_, []), e); meta = smeta }) rhs |> Option.to_list in - if Utils.is_user_ident decl_id then + if Utils.is_user_ident decl_id + then ( let constrain_checks = match dconstrain with | Some Constrain | Some Unconstrain -> - check_transform_shape decl_id decl_var smeta transform - @ constrain_decl dt dconstrain transform decl_id decl_var smeta + check_transform_shape decl_id decl_var smeta transform + @ constrain_decl dt dconstrain transform decl_id decl_var smeta | Some Check -> - check_transform_shape decl_id decl_var smeta transform - @ check_decl decl_var dt decl_id transform smeta dadlevel + check_transform_shape decl_id decl_var smeta transform + @ check_decl decl_var dt decl_id transform smeta dadlevel | None -> [] in - (decl :: rhs_assignment) @ constrain_checks + (decl :: rhs_assignment) @ constrain_checks) else decl :: rhs_assignment +;; let unwrap_block_or_skip = function - | [({Stmt.Fixed.pattern= Block _; _} as b)] | [({pattern= Skip; _} as b)] -> - b - | x -> - raise_s - [%message "Expecting a block or skip, not" (x : Stmt.Located.t list)] + | [ ({ Stmt.Fixed.pattern = Block _; _ } as b) ] | [ ({ pattern = Skip; _ } as b) ] -> b + | x -> raise_s [%message "Expecting a block or skip, not" (x : Stmt.Located.t list)] +;; let dist_name_suffix udf_names name = let is_udf_name s = List.exists ~f:(fun (n, _) -> n = s) udf_names in @@ -475,206 +526,203 @@ let dist_name_suffix udf_names name = Middle.Utils.distribution_suffices |> List.filter ~f:(fun sfx -> Stan_math_signatures.is_stan_math_function_name (name ^ sfx) - || is_udf_name (name ^ sfx) ) + || is_udf_name (name ^ sfx)) |> List.hd with | Some hd -> hd | None -> raise_s [%message "Couldn't find distribution " name] +;; let%expect_test "dist name suffix" = - dist_name_suffix [] "normal" |> print_endline ; + dist_name_suffix [] "normal" |> print_endline; [%expect {| _lpdf |}] +;; let rec trans_stmt ud_dists (declc : decl_context) (ts : Ast.typed_statement) = - let stmt_typed = ts.stmt and smeta = ts.smeta.loc in - let trans_stmt = trans_stmt ud_dists {declc with dconstrain= None} in + let stmt_typed = ts.stmt + and smeta = ts.smeta.loc in + let trans_stmt = trans_stmt ud_dists { declc with dconstrain = None } in let trans_single_stmt s = match trans_stmt s with - | [s] -> s - | s -> Stmt.Fixed.{pattern= SList s; meta= smeta} + | [ s ] -> s + | s -> Stmt.Fixed.{ pattern = SList s; meta = smeta } in - let swrap pattern = [Stmt.Fixed.{meta= smeta; pattern}] in + let swrap pattern = [ Stmt.Fixed.{ meta = smeta; pattern } ] in let mloc = smeta in match stmt_typed with - | Ast.Assignment {assign_lhs; assign_rhs; assign_op} -> - let rec get_lhs_base = function - | {Ast.lval= Ast.LIndexed (l, _); _} -> get_lhs_base l - | {lval= LVariable s; lmeta} -> (s, lmeta) - in - let assign_identifier, lmeta = get_lhs_base assign_lhs in - let id_ad_level = lmeta.Ast.ad_level in - let id_type_ = lmeta.Ast.type_ in - let lhs_type_ = assign_lhs.Ast.lmeta.type_ in - let lhs_ad_level = assign_lhs.Ast.lmeta.ad_level in - let rec get_lhs_indices = function - | {Ast.lval= Ast.LIndexed (l, i); _} -> get_lhs_indices l @ i - | {Ast.lval= Ast.LVariable _; _} -> [] - in - let assign_indices = get_lhs_indices assign_lhs in - let assignee = - { Ast.expr= - ( match assign_indices with - | [] -> Ast.Variable assign_identifier - | _ -> - Ast.Indexed - ( { expr= Ast.Variable assign_identifier - ; emeta= - { Ast.loc= Location_span.empty - ; ad_level= id_ad_level - ; type_= id_type_ } } - , assign_indices ) ) - ; emeta= - { Ast.loc= assign_lhs.lmeta.loc - ; ad_level= lhs_ad_level - ; type_= lhs_type_ } } - in - let rhs = - match assign_op with - | Ast.Assign | Ast.ArrowAssign -> trans_expr assign_rhs - | Ast.OperatorAssign op -> op_to_funapp op [assignee; assign_rhs] - in - Assignment - ( ( assign_identifier.Ast.name - , id_type_ - , List.map ~f:trans_idx assign_indices ) - , rhs ) - |> swrap - | Ast.NRFunApp (fn_kind, {name; _}, args) -> - NRFunApp (trans_fn_kind fn_kind, name, trans_exprs args) |> swrap + | Ast.Assignment { assign_lhs; assign_rhs; assign_op } -> + let rec get_lhs_base = function + | { Ast.lval = Ast.LIndexed (l, _); _ } -> get_lhs_base l + | { lval = LVariable s; lmeta } -> s, lmeta + in + let assign_identifier, lmeta = get_lhs_base assign_lhs in + let id_ad_level = lmeta.Ast.ad_level in + let id_type_ = lmeta.Ast.type_ in + let lhs_type_ = assign_lhs.Ast.lmeta.type_ in + let lhs_ad_level = assign_lhs.Ast.lmeta.ad_level in + let rec get_lhs_indices = function + | { Ast.lval = Ast.LIndexed (l, i); _ } -> get_lhs_indices l @ i + | { Ast.lval = Ast.LVariable _; _ } -> [] + in + let assign_indices = get_lhs_indices assign_lhs in + let assignee = + { Ast.expr = + (match assign_indices with + | [] -> Ast.Variable assign_identifier + | _ -> + Ast.Indexed + ( { expr = Ast.Variable assign_identifier + ; emeta = + { Ast.loc = Location_span.empty + ; ad_level = id_ad_level + ; type_ = id_type_ + } + } + , assign_indices )) + ; emeta = + { Ast.loc = assign_lhs.lmeta.loc; ad_level = lhs_ad_level; type_ = lhs_type_ } + } + in + let rhs = + match assign_op with + | Ast.Assign | Ast.ArrowAssign -> trans_expr assign_rhs + | Ast.OperatorAssign op -> op_to_funapp op [ assignee; assign_rhs ] + in + Assignment + ((assign_identifier.Ast.name, id_type_, List.map ~f:trans_idx assign_indices), rhs) + |> swrap + | Ast.NRFunApp (fn_kind, { name; _ }, args) -> + NRFunApp (trans_fn_kind fn_kind, name, trans_exprs args) |> swrap | Ast.IncrementLogProb e | Ast.TargetPE e -> TargetPE (trans_expr e) |> swrap - | Ast.Tilde {arg; distribution; args; truncation} -> - let suffix = dist_name_suffix ud_dists distribution.name in - let kind = - let possible_names = - List.map ~f:(( ^ ) distribution.name) Utils.distribution_suffices - |> String.Set.of_list - in - if List.exists ~f:(fun (n, _) -> Set.mem possible_names n) ud_dists - then Fun_kind.UserDefined - else StanLib + | Ast.Tilde { arg; distribution; args; truncation } -> + let suffix = dist_name_suffix ud_dists distribution.name in + let kind = + let possible_names = + List.map ~f:(( ^ ) distribution.name) Utils.distribution_suffices + |> String.Set.of_list in - let name = - distribution.name ^ Utils.proportional_to_distribution_infix ^ suffix - in - let add_dist = - Stmt.Fixed.Pattern.TargetPE - Expr. - { Fixed.pattern= FunApp (kind, name, trans_exprs (arg :: args)) - ; meta= - Typed.Meta.create ~type_:UReal ~loc:mloc - ~adlevel:(Ast.expr_ad_lub (arg :: args)) - () } - in - truncate_dist ud_dists distribution arg args truncation @ swrap add_dist + if List.exists ~f:(fun (n, _) -> Set.mem possible_names n) ud_dists + then Fun_kind.UserDefined + else StanLib + in + let name = distribution.name ^ Utils.proportional_to_distribution_infix ^ suffix in + let add_dist = + Stmt.Fixed.Pattern.TargetPE + Expr. + { Fixed.pattern = FunApp (kind, name, trans_exprs (arg :: args)) + ; meta = + Typed.Meta.create + ~type_:UReal + ~loc:mloc + ~adlevel:(Ast.expr_ad_lub (arg :: args)) + () + } + in + truncate_dist ud_dists distribution arg args truncation @ swrap add_dist | Ast.Print ps -> - NRFunApp - ( CompilerInternal - , Internal_fun.to_string FnPrint - , trans_printables smeta ps ) - |> swrap + NRFunApp (CompilerInternal, Internal_fun.to_string FnPrint, trans_printables smeta ps) + |> swrap | Ast.Reject ps -> - NRFunApp - ( CompilerInternal - , Internal_fun.to_string FnReject - , trans_printables smeta ps ) - |> swrap + NRFunApp (CompilerInternal, Internal_fun.to_string FnReject, trans_printables smeta ps) + |> swrap | Ast.IfThenElse (cond, ifb, elseb) -> - IfElse - ( trans_expr cond - , trans_single_stmt ifb - , Option.map ~f:trans_single_stmt elseb ) - |> swrap - | Ast.While (cond, body) -> - While (trans_expr cond, trans_single_stmt body) |> swrap - | Ast.For {loop_variable; lower_bound; upper_bound; loop_body} -> - let body = - match trans_single_stmt loop_body with - | {pattern= Block _; _} as b -> b - | x -> {x with pattern= Block [x]} - in - For - { loopvar= loop_variable.Ast.name - ; lower= trans_expr lower_bound - ; upper= trans_expr upper_bound - ; body } - |> swrap + IfElse (trans_expr cond, trans_single_stmt ifb, Option.map ~f:trans_single_stmt elseb) + |> swrap + | Ast.While (cond, body) -> While (trans_expr cond, trans_single_stmt body) |> swrap + | Ast.For { loop_variable; lower_bound; upper_bound; loop_body } -> + let body = + match trans_single_stmt loop_body with + | { pattern = Block _; _ } as b -> b + | x -> { x with pattern = Block [ x ] } + in + For + { loopvar = loop_variable.Ast.name + ; lower = trans_expr lower_bound + ; upper = trans_expr upper_bound + ; body + } + |> swrap | Ast.ForEach (loopvar, iteratee, body) -> - let iteratee' = trans_expr iteratee in - let body_stmts = - match trans_single_stmt body with - | {pattern= Block body_stmts; _} -> body_stmts - | b -> [b] - in - let decl_type = - match Expr.Typed.type_of iteratee' with - | UMatrix -> UnsizedType.UReal - | t -> - Expr.Helpers.(infer_type_of_indexed t [Index.Single loop_bottom]) - in - let decl_loopvar = - Stmt.Fixed. - { meta= smeta - ; pattern= - Decl - { decl_adtype= Expr.Typed.adlevel_of iteratee' - ; decl_id= loopvar.name - ; decl_type= Unsized decl_type } } - in - let assignment var = - Stmt.Fixed. - { pattern= Assignment ((loopvar.name, decl_type, []), var) - ; meta= smeta } - in - let bodyfn var = - Stmt.Fixed. - { pattern= Block (decl_loopvar :: assignment var :: body_stmts) - ; meta= smeta } - in - Stmt.Helpers.[ensure_var (for_each bodyfn) iteratee' smeta] + let iteratee' = trans_expr iteratee in + let body_stmts = + match trans_single_stmt body with + | { pattern = Block body_stmts; _ } -> body_stmts + | b -> [ b ] + in + let decl_type = + match Expr.Typed.type_of iteratee' with + | UMatrix -> UnsizedType.UReal + | t -> Expr.Helpers.(infer_type_of_indexed t [ Index.Single loop_bottom ]) + in + let decl_loopvar = + Stmt.Fixed. + { meta = smeta + ; pattern = + Decl + { decl_adtype = Expr.Typed.adlevel_of iteratee' + ; decl_id = loopvar.name + ; decl_type = Unsized decl_type + } + } + in + let assignment var = + Stmt.Fixed. + { pattern = Assignment ((loopvar.name, decl_type, []), var); meta = smeta } + in + let bodyfn var = + Stmt.Fixed. + { pattern = Block (decl_loopvar :: assignment var :: body_stmts); meta = smeta } + in + Stmt.Helpers.[ ensure_var (for_each bodyfn) iteratee' smeta ] | Ast.FunDef _ -> - raise_s - [%message - "Found function definition statement outside of function block"] - | Ast.VarDecl - {decl_type; transformation; identifier; initial_value; is_global= _} -> - trans_decl declc smeta decl_type - (Program.map_transformation trans_expr transformation) - identifier initial_value + raise_s [%message "Found function definition statement outside of function block"] + | Ast.VarDecl { decl_type; transformation; identifier; initial_value; is_global = _ } -> + trans_decl + declc + smeta + decl_type + (Program.map_transformation trans_expr transformation) + identifier + initial_value | Ast.Block stmts -> Block (List.concat_map ~f:trans_stmt stmts) |> swrap | Ast.Return e -> Return (Some (trans_expr e)) |> swrap | Ast.ReturnVoid -> Return None |> swrap | Ast.Break -> Break |> swrap | Ast.Continue -> Continue |> swrap | Ast.Skip -> Skip |> swrap +;; let trans_fun_def ud_dists (ts : Ast.typed_statement) = match ts.stmt with - | Ast.FunDef {returntype; funname; arguments; body} -> - [ Program. - { fdrt= - (match returntype with Void -> None | ReturnType ut -> Some ut) - ; fdname= funname.name - ; fdargs= List.map ~f:trans_arg arguments - ; fdbody= - trans_stmt ud_dists - {dconstrain= None; dadlevel= AutoDiffable} - body - |> unwrap_block_or_skip - ; fdloc= ts.smeta.loc } ] - | _ -> - raise_s - [%message "Found non-function definition statement in function block"] + | Ast.FunDef { returntype; funname; arguments; body } -> + [ Program. + { fdrt = + (match returntype with + | Void -> None + | ReturnType ut -> Some ut) + ; fdname = funname.name + ; fdargs = List.map ~f:trans_arg arguments + ; fdbody = + trans_stmt ud_dists { dconstrain = None; dadlevel = AutoDiffable } body + |> unwrap_block_or_skip + ; fdloc = ts.smeta.loc + } + ] + | _ -> raise_s [%message "Found non-function definition statement in function block"] +;; let get_block block prog = match block with | Program.Parameters -> prog.Ast.parametersblock | TransformedParameters -> prog.transformedparametersblock | GeneratedQuantities -> prog.generatedquantitiesblock +;; let migrate_checks_to_end_of_block stmts = let is_check = Stmt.Helpers.contains_fn FnCheck in let checks, not_checks = List.partition_tf ~f:is_check stmts in not_checks @ checks +;; let trans_prog filename (p : Ast.typed_program) : Program.Typed.t = let { Ast.functionblock @@ -682,21 +730,24 @@ let trans_prog filename (p : Ast.typed_program) : Program.Typed.t = ; transformeddatablock ; parametersblock ; transformedparametersblock - ; modelblock; _ } = + ; modelblock + ; _ + } + = p in let map f list_op = Option.value ~default:[] list_op |> List.concat_map ~f in let grab_fundef_names_and_types = function - | {Ast.stmt= Ast.FunDef {funname; arguments= (_, type_, _) :: _; _}; _} -> - [(funname.name, type_)] + | { Ast.stmt = Ast.FunDef { funname; arguments = (_, type_, _) :: _; _ }; _ } -> + [ funname.name, type_ ] | _ -> [] in let ud_dists = map grab_fundef_names_and_types functionblock in let trans_stmt = trans_stmt ud_dists in let get_name_size s = match s.Ast.stmt with - | Ast.VarDecl {decl_type= Sized st; identifier; transformation; _} -> - [(identifier.name, trans_sizedtype st, transformation)] + | Ast.VarDecl { decl_type = Sized st; identifier; transformation; _ } -> + [ identifier.name, trans_sizedtype st, transformation ] | _ -> [] in let grab_names_sizes block = @@ -706,89 +757,90 @@ let trans_prog filename (p : Ast.typed_program) : Program.Typed.t = (List.map ~f:(fun (n, s, t) -> ( n , Program. - { out_constrained_st= s - ; out_unconstrained_st= param_size t s - ; out_block= block - ; out_trans= map_transformation trans_expr t } ) )) + { out_constrained_st = s + ; out_unconstrained_st = param_size t s + ; out_block = block + ; out_trans = map_transformation trans_expr t + } ))) in let output_vars = grab_names_sizes Parameters @ grab_names_sizes TransformedParameters @ grab_names_sizes GeneratedQuantities - and input_vars = - map get_name_size datablock |> List.map ~f:(fun (n, st, _) -> (n, st)) - in - let declc = {dconstrain= None; dadlevel= DataOnly} in + and input_vars = map get_name_size datablock |> List.map ~f:(fun (n, st, _) -> n, st) in + let declc = { dconstrain = None; dadlevel = DataOnly } in let datab = - map (trans_stmt {declc with dconstrain= Some Check}) datablock + map (trans_stmt { declc with dconstrain = Some Check }) datablock |> migrate_checks_to_end_of_block in let prepare_data = - datab - @ map (trans_stmt {declc with dconstrain= Some Check}) transformeddatablock + datab @ map (trans_stmt { declc with dconstrain = Some Check }) transformeddatablock |> migrate_checks_to_end_of_block in - let modelb = - map (trans_stmt {declc with dadlevel= AutoDiffable}) modelblock - in + let modelb = map (trans_stmt { declc with dadlevel = AutoDiffable }) modelblock in let log_prob = map - (trans_stmt {dconstrain= Some Constrain; dadlevel= AutoDiffable}) + (trans_stmt { dconstrain = Some Constrain; dadlevel = AutoDiffable }) parametersblock - @ ( map - (trans_stmt {dconstrain= Some Check; dadlevel= AutoDiffable}) - transformedparametersblock - |> migrate_checks_to_end_of_block ) + @ (map + (trans_stmt { dconstrain = Some Check; dadlevel = AutoDiffable }) + transformedparametersblock + |> migrate_checks_to_end_of_block) @ match modelb with | [] -> [] - | hd :: _ -> [{pattern= Block modelb; meta= hd.meta}] - in - let gen_from_block declc block = - map (trans_stmt declc) (get_block block p) + | hd :: _ -> [ { pattern = Block modelb; meta = hd.meta } ] in + let gen_from_block declc block = map (trans_stmt declc) (get_block block p) in let txparam_decls, txparam_stmts = gen_from_block declc TransformedParameters |> List.partition_tf ~f:(function - | {pattern= Decl _; _} -> true - | _ -> false ) + | { pattern = Decl _; _ } -> true + | _ -> false) in let compiler_if_return cond = Stmt.Fixed. - { pattern= - IfElse (cond, {pattern= Return None; meta= Location_span.empty}, None) - ; meta= Location_span.empty } + { pattern = + IfElse (cond, { pattern = Return None; meta = Location_span.empty }, None) + ; meta = Location_span.empty + } in - let iexpr pattern = Expr.{pattern; Fixed.meta= Typed.Meta.empty} in - let fnot e = FunApp (StanLib, Operator.to_string PNot, [e]) |> iexpr in + let iexpr pattern = Expr.{ pattern; Fixed.meta = Typed.Meta.empty } in + let fnot e = FunApp (StanLib, Operator.to_string PNot, [ e ]) |> iexpr in let tparam_early_return = let to_var fv = iexpr (Var (Flag_vars.to_string fv)) in let v1 = to_var EmitTransformedParameters in let v2 = to_var EmitGeneratedQuantities in - [compiler_if_return (fnot (EOr (v1, v2) |> iexpr))] + [ compiler_if_return (fnot (EOr (v1, v2) |> iexpr)) ] in let gq_stmts = migrate_checks_to_end_of_block - (gen_from_block {declc with dconstrain= Some Check} GeneratedQuantities) + (gen_from_block { declc with dconstrain = Some Check } GeneratedQuantities) in let gq_early_return = [ compiler_if_return - (fnot (Var (Flag_vars.to_string EmitGeneratedQuantities) |> iexpr)) ] + (fnot (Var (Flag_vars.to_string EmitGeneratedQuantities) |> iexpr)) + ] in let generate_quantities = - gen_from_block {declc with dconstrain= Some Constrain} Parameters - @ txparam_decls @ tparam_early_return @ txparam_stmts @ gq_early_return + gen_from_block { declc with dconstrain = Some Constrain } Parameters + @ txparam_decls + @ tparam_early_return + @ txparam_stmts + @ gq_early_return @ gq_stmts in let transform_inits = - gen_from_block {declc with dconstrain= Some Unconstrain} Parameters + gen_from_block { declc with dconstrain = Some Unconstrain } Parameters in - { functions_block= map (trans_fun_def ud_dists) functionblock + { functions_block = map (trans_fun_def ud_dists) functionblock ; input_vars ; prepare_data ; log_prob ; generate_quantities ; transform_inits ; output_vars - ; prog_name= !Semantic_check.model_name - ; prog_path= filename } + ; prog_name = !Semantic_check.model_name + ; prog_path = filename + } +;; diff --git a/src/frontend/Canonicalize.ml b/src/frontend/Canonicalize.ml index 7772b61e59..174258b207 100644 --- a/src/frontend/Canonicalize.ml +++ b/src/frontend/Canonicalize.ml @@ -3,241 +3,257 @@ open Ast let deprecated_functions = String.Map.of_alist_exn - [ ("multiply_log", "lmultiply") - ; ("binomial_coefficient_log", "lchoose") - ; ("integrate_ode", "integrate_ode_rk45") ] + [ "multiply_log", "lmultiply" + ; "binomial_coefficient_log", "lchoose" + ; "integrate_ode", "integrate_ode_rk45" + ] +;; let deprecated_distributions = String.Map.of_alist_exn - (List.concat_map Middle.Stan_math_signatures.distributions + (List.concat_map + Middle.Stan_math_signatures.distributions ~f:(fun (fnkinds, name, _) -> List.filter_map fnkinds ~f:(function - | Lpdf -> Some (name ^ "_log", name ^ "_lpdf") - | Lpmf -> Some (name ^ "_log", name ^ "_lpmf") - | Cdf -> Some (name ^ "_cdf_log", name ^ "_lcdf") - | Ccdf -> Some (name ^ "_ccdf_log", name ^ "_lccdf") - | Rng | UnaryVectorized -> None ) )) + | Lpdf -> Some (name ^ "_log", name ^ "_lpdf") + | Lpmf -> Some (name ^ "_log", name ^ "_lpmf") + | Cdf -> Some (name ^ "_cdf_log", name ^ "_lcdf") + | Ccdf -> Some (name ^ "_ccdf_log", name ^ "_lccdf") + | Rng | UnaryVectorized -> None))) +;; let deprecated_userdefined = String.Table.create () - -let is_distribution name = - Option.is_some (String.Map.find deprecated_distributions name) +let is_distribution name = Option.is_some (String.Map.find deprecated_distributions name) let rename_distribution name = Option.value ~default:name (String.Map.find deprecated_distributions name) +;; let rename_function name = Option.value ~default:name (String.Map.find deprecated_functions name) +;; let distribution_suffix name = String.is_suffix ~suffix:"_lpdf" name || String.is_suffix ~suffix:"_lpmf" name || String.is_suffix ~suffix:"_lcdf" name || String.is_suffix ~suffix:"_lccdf" name +;; let userdef_distributions stmts = List.filter_map ~f:(function - | {stmt= FunDef {funname= {name; _}; _}; _} -> - if - String.is_suffix ~suffix:"_log_lpdf" name - || String.is_suffix ~suffix:"_log_lpmf" name - then Some (String.drop_suffix name 5) - else if String.is_suffix ~suffix:"_log_log" name then - Some (String.drop_suffix name 4) - else None + | { stmt = FunDef { funname = { name; _ }; _ }; _ } -> + if String.is_suffix ~suffix:"_log_lpdf" name + || String.is_suffix ~suffix:"_log_lpmf" name + then Some (String.drop_suffix name 5) + else if String.is_suffix ~suffix:"_log_log" name + then Some (String.drop_suffix name 4) + else None | _ -> None) (Option.value ~default:[] stmts) +;; let without_suffix user_dists name = - if - String.is_suffix ~suffix:"_lpdf" name - || String.is_suffix ~suffix:"_lpmf" name + if String.is_suffix ~suffix:"_lpdf" name || String.is_suffix ~suffix:"_lpmf" name then String.drop_suffix name 5 - else if - String.is_suffix ~suffix:"_log" name - && not - ( is_distribution (name ^ "_log") - || List.exists ~f:(( = ) name) user_dists ) + else if String.is_suffix ~suffix:"_log" name + && not + (is_distribution (name ^ "_log") || List.exists ~f:(( = ) name) user_dists) then String.drop_suffix name 4 else name +;; -let rec repair_syntax_expr {expr; emeta} = +let rec repair_syntax_expr { expr; emeta } = let expr = match expr with - | FunApp (f, {name; id_loc}, e) when distribution_suffix name -> - CondDistApp (f, {name; id_loc}, List.map ~f:repair_syntax_expr e) - | CondDistApp (f, {name; id_loc}, e) when not (distribution_suffix name) -> - FunApp (f, {name; id_loc}, List.map ~f:repair_syntax_expr e) + | FunApp (f, { name; id_loc }, e) when distribution_suffix name -> + CondDistApp (f, { name; id_loc }, List.map ~f:repair_syntax_expr e) + | CondDistApp (f, { name; id_loc }, e) when not (distribution_suffix name) -> + FunApp (f, { name; id_loc }, List.map ~f:repair_syntax_expr e) | _ -> map_expression repair_syntax_expr ident expr in - {expr; emeta} + { expr; emeta } +;; let repair_syntax_lval = map_lval_with repair_syntax_expr ident -let rec repair_syntax_stmt user_dists {stmt; smeta} = +let rec repair_syntax_stmt user_dists { stmt; smeta } = match stmt with - | Tilde {arg; distribution= {name; id_loc}; args; truncation} -> - { stmt= - Tilde - { arg= repair_syntax_expr arg - ; distribution= {name= without_suffix user_dists name; id_loc} - ; args= List.map ~f:repair_syntax_expr args - ; truncation= map_truncation repair_syntax_expr truncation } - ; smeta } + | Tilde { arg; distribution = { name; id_loc }; args; truncation } -> + { stmt = + Tilde + { arg = repair_syntax_expr arg + ; distribution = { name = without_suffix user_dists name; id_loc } + ; args = List.map ~f:repair_syntax_expr args + ; truncation = map_truncation repair_syntax_expr truncation + } + ; smeta + } | _ -> - { stmt= - map_statement repair_syntax_expr - (repair_syntax_stmt user_dists) - repair_syntax_lval ident stmt - ; smeta } + { stmt = + map_statement + repair_syntax_expr + (repair_syntax_stmt user_dists) + repair_syntax_lval + ident + stmt + ; smeta + } +;; -let rec replace_deprecated_expr {expr; emeta} = +let rec replace_deprecated_expr { expr; emeta } = let expr = match expr with | GetLP -> GetTarget - | FunApp (StanLib, {name= "abs"; id_loc}, [e]) + | FunApp (StanLib, { name = "abs"; id_loc }, [ e ]) when Middle.UnsizedType.is_real_type e.emeta.type_ -> - FunApp (StanLib, {name= "fabs"; id_loc}, [replace_deprecated_expr e]) - | FunApp (StanLib, {name= "if_else"; _}, [c; t; e]) -> - Paren - (replace_deprecated_expr - {expr= TernaryIf ({expr= Paren c; emeta= c.emeta}, t, e); emeta}) - | FunApp (StanLib, {name; id_loc}, e) -> - if is_distribution name then - CondDistApp - ( StanLib - , {name= rename_distribution name; id_loc} - , List.map ~f:replace_deprecated_expr e ) - else - FunApp - ( StanLib - , {name= rename_function name; id_loc} - , List.map ~f:replace_deprecated_expr e ) - | FunApp (UserDefined, {name; id_loc}, e) -> ( - match String.Table.find deprecated_userdefined name with + FunApp (StanLib, { name = "fabs"; id_loc }, [ replace_deprecated_expr e ]) + | FunApp (StanLib, { name = "if_else"; _ }, [ c; t; e ]) -> + Paren + (replace_deprecated_expr + { expr = TernaryIf ({ expr = Paren c; emeta = c.emeta }, t, e); emeta }) + | FunApp (StanLib, { name; id_loc }, e) -> + if is_distribution name + then + CondDistApp + ( StanLib + , { name = rename_distribution name; id_loc } + , List.map ~f:replace_deprecated_expr e ) + else + FunApp + ( StanLib + , { name = rename_function name; id_loc } + , List.map ~f:replace_deprecated_expr e ) + | FunApp (UserDefined, { name; id_loc }, e) -> + (match String.Table.find deprecated_userdefined name with | Some newname -> - CondDistApp - ( UserDefined - , {name= newname; id_loc} - , List.map ~f:replace_deprecated_expr e ) + CondDistApp + (UserDefined, { name = newname; id_loc }, List.map ~f:replace_deprecated_expr e) | None -> - FunApp - (UserDefined, {name; id_loc}, List.map ~f:replace_deprecated_expr e) - ) + FunApp (UserDefined, { name; id_loc }, List.map ~f:replace_deprecated_expr e)) | _ -> map_expression replace_deprecated_expr ident expr in - {expr; emeta} + { expr; emeta } +;; let replace_deprecated_lval = map_lval_with replace_deprecated_expr ident -let rec replace_deprecated_stmt {stmt; smeta} = +let rec replace_deprecated_stmt { stmt; smeta } = let stmt = match stmt with | IncrementLogProb e -> TargetPE (replace_deprecated_expr e) - | Assignment {assign_lhs= l; assign_op= ArrowAssign; assign_rhs= e} -> - Assignment - { assign_lhs= replace_deprecated_lval l - ; assign_op= Assign - ; assign_rhs= replace_deprecated_expr e } - | FunDef {returntype; funname= {name; id_loc}; arguments; body} -> - FunDef - { returntype - ; funname= - { name= - Option.value ~default:name - (String.Table.find deprecated_userdefined name) - ; id_loc } - ; arguments - ; body= replace_deprecated_stmt body } + | Assignment { assign_lhs = l; assign_op = ArrowAssign; assign_rhs = e } -> + Assignment + { assign_lhs = replace_deprecated_lval l + ; assign_op = Assign + ; assign_rhs = replace_deprecated_expr e + } + | FunDef { returntype; funname = { name; id_loc }; arguments; body } -> + FunDef + { returntype + ; funname = + { name = + Option.value ~default:name (String.Table.find deprecated_userdefined name) + ; id_loc + } + ; arguments + ; body = replace_deprecated_stmt body + } | _ -> - map_statement replace_deprecated_expr replace_deprecated_stmt - replace_deprecated_lval ident stmt + map_statement + replace_deprecated_expr + replace_deprecated_stmt + replace_deprecated_lval + ident + stmt in - {stmt; smeta} + { stmt; smeta } +;; -let rec no_parens {expr; emeta} = +let rec no_parens { expr; emeta } = match expr with | Paren e -> no_parens e - | Variable _ | IntNumeral _ | RealNumeral _ | GetLP | GetTarget -> - {expr; emeta} + | Variable _ | IntNumeral _ | RealNumeral _ | GetLP | GetTarget -> { expr; emeta } | TernaryIf _ | BinOp _ | PrefixOp _ | PostfixOp _ -> - {expr= map_expression keep_parens ident expr; emeta} + { expr = map_expression keep_parens ident expr; emeta } | Indexed (e, l) -> - { expr= - Indexed - ( keep_parens e - , List.map - ~f:(function - | Single e -> Single (no_parens e) - | i -> map_index keep_parens i) - l ) - ; emeta } + { expr = + Indexed + ( keep_parens e + , List.map + ~f:(function + | Single e -> Single (no_parens e) + | i -> map_index keep_parens i) + l ) + ; emeta + } | ArrayExpr _ | RowVectorExpr _ | FunApp _ | CondDistApp _ -> - {expr= map_expression no_parens ident expr; emeta} + { expr = map_expression no_parens ident expr; emeta } -and keep_parens {expr; emeta} = +and keep_parens { expr; emeta } = match expr with - | Paren {expr= Paren e; _} -> keep_parens e - | Paren ({expr= BinOp _; _} as e) - |Paren ({expr= PrefixOp _; _} as e) - |Paren ({expr= PostfixOp _; _} as e) - |Paren ({expr= TernaryIf _; _} as e) -> - {expr= Paren (no_parens e); emeta} - | _ -> no_parens {expr; emeta} + | Paren { expr = Paren e; _ } -> keep_parens e + | Paren ({ expr = BinOp _; _ } as e) + | Paren ({ expr = PrefixOp _; _ } as e) + | Paren ({ expr = PostfixOp _; _ } as e) + | Paren ({ expr = TernaryIf _; _ } as e) -> { expr = Paren (no_parens e); emeta } + | _ -> no_parens { expr; emeta } +;; let parens_lval = map_lval_with no_parens ident -let rec parens_stmt {stmt; smeta} = +let rec parens_stmt { stmt; smeta } = let stmt = match stmt with | VarDecl - { decl_type= d - ; transformation= t + { decl_type = d; transformation = t; identifier; initial_value = init; is_global } + -> + VarDecl + { decl_type = Middle.Type.map no_parens d + ; transformation = Middle.Program.map_transformation keep_parens t ; identifier - ; initial_value= init - ; is_global } -> - VarDecl - { decl_type= Middle.Type.map no_parens d - ; transformation= Middle.Program.map_transformation keep_parens t - ; identifier - ; initial_value= Option.map ~f:no_parens init - ; is_global } - | For {loop_variable; lower_bound; upper_bound; loop_body} -> - For - { loop_variable - ; lower_bound= keep_parens lower_bound - ; upper_bound= keep_parens upper_bound - ; loop_body= parens_stmt loop_body } + ; initial_value = Option.map ~f:no_parens init + ; is_global + } + | For { loop_variable; lower_bound; upper_bound; loop_body } -> + For + { loop_variable + ; lower_bound = keep_parens lower_bound + ; upper_bound = keep_parens upper_bound + ; loop_body = parens_stmt loop_body + } | _ -> map_statement no_parens parens_stmt parens_lval ident stmt in - {stmt; smeta} + { stmt; smeta } +;; let repair_syntax program : untyped_program = - map_program - (repair_syntax_stmt (userdef_distributions program.functionblock)) - program + map_program (repair_syntax_stmt (userdef_distributions program.functionblock)) program +;; let canonicalize_program program : typed_program = - String.Table.clear deprecated_userdefined ; + String.Table.clear deprecated_userdefined; program.functionblock |> Option.iter ~f: (List.iter ~f:(function - | { stmt= - FunDef {funname= {name; _}; arguments= (_, type_, _) :: _; _} - ; smeta= _ } - when String.is_suffix ~suffix:"_log" name -> + | { stmt = + FunDef { funname = { name; _ }; arguments = (_, type_, _) :: _; _ } + ; smeta = _ + } + when String.is_suffix ~suffix:"_log" name -> let newname = - if String.is_suffix ~suffix:"_cdf_log" name then - String.drop_suffix name 8 ^ "_lcdf" - else if String.is_suffix ~suffix:"_ccdf_log" name then - String.drop_suffix name 9 ^ "_lccdf" - else if Middle.UnsizedType.is_real_type type_ then - String.drop_suffix name 4 ^ "_lpdf" + if String.is_suffix ~suffix:"_cdf_log" name + then String.drop_suffix name 8 ^ "_lcdf" + else if String.is_suffix ~suffix:"_ccdf_log" name + then String.drop_suffix name 9 ^ "_lccdf" + else if Middle.UnsizedType.is_real_type type_ + then String.drop_suffix name 4 ^ "_lpdf" else String.drop_suffix name 4 ^ "_lpmf" in String.Table.add deprecated_userdefined ~key:name ~data:newname - |> (ignore : [`Ok | `Duplicate] -> unit) - | _ -> () )) ; + |> (ignore : [ `Ok | `Duplicate ] -> unit) + | _ -> ())); program |> map_program replace_deprecated_stmt |> map_program parens_stmt +;; diff --git a/src/frontend/Debug_data_generation.ml b/src/frontend/Debug_data_generation.ml index db02f5ce79..43783ed2d7 100644 --- a/src/frontend/Debug_data_generation.ml +++ b/src/frontend/Debug_data_generation.ml @@ -5,27 +5,32 @@ open Ast let rec transpose = function | [] :: _ -> [] | rows -> - let hd = List.map ~f:List.hd_exn rows in - let tl = List.map ~f:List.tl_exn rows in - hd :: transpose tl + let hd = List.map ~f:List.hd_exn rows in + let tl = List.map ~f:List.tl_exn rows in + hd :: transpose tl +;; let dotproduct xs ys = List.fold2_exn xs ys ~init:0. ~f:(fun accum x y -> accum +. (x *. y)) +;; let matprod x y = let y_T = transpose y in - if List.length x <> List.length y_T then - failwith "Matrix multiplication dim. mismatch" + if List.length x <> List.length y_T + then failwith "Matrix multiplication dim. mismatch" else List.map ~f:(fun row -> List.map ~f:(dotproduct row) y_T) x +;; let rec vect_to_mat l m = let len = List.length l in - if len % m <> 0 then - failwith "the length has to be a whole multiple of the partition size" - else if len = m then [l] - else + if len % m <> 0 + then failwith "the length has to be a whole multiple of the partition size" + else if len = m + then [ l ] + else ( let hd, tl = List.split_n l m in - hd :: vect_to_mat tl m + hd :: vect_to_mat tl m) +;; let unwrap_num_exn m e = let e = Ast_to_Mir.trans_expr e in @@ -35,111 +40,129 @@ let unwrap_num_exn m e = match e.pattern with | Lit (_, s) -> Float.of_string s | _ -> raise_s [%sexp ("Cannot convert size to number." : string)] +;; let unwrap_int_exn m e = Int.of_float (unwrap_num_exn m e) let gen_num_int m t = - let def_low, diff = (2, 4) in + let def_low, diff = 2, 4 in let low, up = match t with - | Program.Lower e -> (unwrap_int_exn m e, unwrap_int_exn m e + diff) - | Upper e -> (unwrap_int_exn m e - diff, unwrap_int_exn m e) - | LowerUpper (e1, e2) -> (unwrap_int_exn m e1, unwrap_int_exn m e2) - | _ -> (def_low, def_low + diff) + | Program.Lower e -> unwrap_int_exn m e, unwrap_int_exn m e + diff + | Upper e -> unwrap_int_exn m e - diff, unwrap_int_exn m e + | LowerUpper (e1, e2) -> unwrap_int_exn m e1, unwrap_int_exn m e2 + | _ -> def_low, def_low + diff in let low = if low = 0 && up <> 1 then low + 1 else low in Random.int (up - low + 1) + low +;; let gen_num_real m t = - let def_low, diff = (2., 5.) in + let def_low, diff = 2., 5. in let low, up = match t with - | Program.Lower e -> (unwrap_num_exn m e, unwrap_num_exn m e +. diff) - | Upper e -> (unwrap_num_exn m e -. diff, unwrap_num_exn m e) - | LowerUpper (e1, e2) -> (unwrap_num_exn m e1, unwrap_num_exn m e2) - | _ -> (def_low, def_low +. diff) + | Program.Lower e -> unwrap_num_exn m e, unwrap_num_exn m e +. diff + | Upper e -> unwrap_num_exn m e -. diff, unwrap_num_exn m e + | LowerUpper (e1, e2) -> unwrap_num_exn m e1, unwrap_num_exn m e2 + | _ -> def_low, def_low +. diff in Random.float_range low up +;; let rec repeat n e = - match n with n when n <= 0 -> [] | m -> e :: repeat (m - 1) e + match n with + | n when n <= 0 -> [] + | m -> e :: repeat (m - 1) e +;; let rec repeat_th n f = - match n with n when n <= 0 -> [] | m -> f () :: repeat_th (m - 1) f + match n with + | n when n <= 0 -> [] + | m -> f () :: repeat_th (m - 1) f +;; let wrap_int n = - { expr= IntNumeral (Int.to_string n) - ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UInt} } + { expr = IntNumeral (Int.to_string n) + ; emeta = { loc = Location_span.empty; ad_level = DataOnly; type_ = UInt } + } +;; let int_two = wrap_int 2 let wrap_real r = - { expr= RealNumeral (Float.to_string r) - ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UReal} } + { expr = RealNumeral (Float.to_string r) + ; emeta = { loc = Location_span.empty; ad_level = DataOnly; type_ = UReal } + } +;; let wrap_row_vector l = - { expr= RowVectorExpr l - ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= URowVector} } + { expr = RowVectorExpr l + ; emeta = { loc = Location_span.empty; ad_level = DataOnly; type_ = URowVector } + } +;; let wrap_vector l = - { expr= PostfixOp (wrap_row_vector l, Transpose) - ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UVector} } + { expr = PostfixOp (wrap_row_vector l, Transpose) + ; emeta = { loc = Location_span.empty; ad_level = DataOnly; type_ = UVector } + } +;; let gen_int m t = wrap_int (gen_num_int m t) let gen_real m t = wrap_real (gen_num_real m t) let gen_row_vector m n t = - { expr= RowVectorExpr (repeat_th n (fun _ -> gen_real m t)) - ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UMatrix} } + { expr = RowVectorExpr (repeat_th n (fun _ -> gen_real m t)) + ; emeta = { loc = Location_span.empty; ad_level = DataOnly; type_ = UMatrix } + } +;; let gen_vector m n t = let gen_ordered n = let l = repeat_th n (fun _ -> Random.float 1.) in let l = - List.fold (List.tl_exn l) ~init:[List.hd_exn l] ~f:(fun accum elt -> - (Float.exp elt +. List.hd_exn accum) :: accum ) + List.fold (List.tl_exn l) ~init:[ List.hd_exn l ] ~f:(fun accum elt -> + (Float.exp elt +. List.hd_exn accum) :: accum) in l in match t with | Program.Simplex -> - let l = repeat_th n (fun _ -> Random.float 1.) in - let sum = List.fold l ~init:0. ~f:(fun accum elt -> accum +. elt) in - let l = List.map l ~f:(fun x -> x /. sum) in - wrap_vector (List.map ~f:wrap_real l) + let l = repeat_th n (fun _ -> Random.float 1.) in + let sum = List.fold l ~init:0. ~f:(fun accum elt -> accum +. elt) in + let l = List.map l ~f:(fun x -> x /. sum) in + wrap_vector (List.map ~f:wrap_real l) | Ordered -> - let l = gen_ordered n in - let halfmax = - Option.value_exn (List.max_elt l ~compare:compare_float) /. 2. - in - let l = List.map l ~f:(fun x -> (x -. halfmax) /. halfmax) in - wrap_vector (List.map ~f:wrap_real l) + let l = gen_ordered n in + let halfmax = Option.value_exn (List.max_elt l ~compare:compare_float) /. 2. in + let l = List.map l ~f:(fun x -> (x -. halfmax) /. halfmax) in + wrap_vector (List.map ~f:wrap_real l) | PositiveOrdered -> - let l = gen_ordered n in - let max = Option.value_exn (List.max_elt l ~compare:compare_float) in - let l = List.map l ~f:(fun x -> x /. max) in - wrap_vector (List.map ~f:wrap_real l) + let l = gen_ordered n in + let max = Option.value_exn (List.max_elt l ~compare:compare_float) in + let l = List.map l ~f:(fun x -> x /. max) in + wrap_vector (List.map ~f:wrap_real l) | UnitVector -> - let l = repeat_th n (fun _ -> Random.float 1.) in - let sum = - Float.sqrt - (List.fold l ~init:0. ~f:(fun accum elt -> accum +. (elt ** 2.))) - in - let l = List.map l ~f:(fun x -> x /. sum) in - wrap_vector (List.map ~f:wrap_real l) - | _ -> {int_two with expr= PostfixOp (gen_row_vector m n t, Transpose)} + let l = repeat_th n (fun _ -> Random.float 1.) in + let sum = + Float.sqrt (List.fold l ~init:0. ~f:(fun accum elt -> accum +. (elt ** 2.))) + in + let l = List.map l ~f:(fun x -> x /. sum) in + wrap_vector (List.map ~f:wrap_real l) + | _ -> { int_two with expr = PostfixOp (gen_row_vector m n t, Transpose) } +;; let gen_cov_unwrapped n = let l = repeat_th (n * n) (fun _ -> Random.float 2.) in let l_mat = vect_to_mat l n in matprod l_mat (transpose l_mat) +;; let wrap_real_mat m = let mat_wrapped = - List.map ~f:wrap_row_vector - (List.map ~f:(fun x -> List.map ~f:wrap_real x) m) + List.map ~f:wrap_row_vector (List.map ~f:(fun x -> List.map ~f:wrap_real x) m) in - {int_two with expr= RowVectorExpr mat_wrapped} + { int_two with expr = RowVectorExpr mat_wrapped } +;; let gen_diag_mat l = let n = List.length l in @@ -147,8 +170,9 @@ let gen_diag_mat l = (List.range 1 (n + 1)) ~f:(fun k -> repeat (min (k - 1) n) 0. - @ (if k <= n then [List.nth_exn l (k - 1)] else []) - @ repeat (n - k) 0. ) + @ (if k <= n then [ List.nth_exn l (k - 1) ] else []) + @ repeat (n - k) 0.) +;; let fill_lower_triangular m = let fill_row i l = @@ -156,17 +180,20 @@ let fill_lower_triangular m = List.init ~f:(fun _ -> Random.float 2.) i @ tl in List.mapi ~f:fill_row m +;; let pad_mat mm m n = let padding_mat = List.init (m - n) ~f:(fun _ -> List.init n ~f:(fun _ -> Random.float 2.)) in wrap_real_mat (mm @ padding_mat) +;; let gen_cov_cholesky m n = let diag_mat = gen_diag_mat (List.init ~f:(fun _ -> Random.float 2.) n) in let filled_mat = fill_lower_triangular diag_mat in if m <= n then wrap_real_mat filled_mat else pad_mat filled_mat m n +;; let gen_corr_cholesky_unwrapped n = let diag_mat = gen_diag_mat (List.init ~f:(fun _ -> Random.float 2.) n) in @@ -178,6 +205,7 @@ let gen_corr_cholesky_unwrapped n = List.map ~f:(fun x -> x /. row_norm) l in List.map ~f:row_normalizer filled_mat +;; let gen_corr_cholesky n = wrap_real_mat (gen_corr_cholesky_unwrapped n) @@ -188,10 +216,12 @@ let gen_corr_cholesky n = wrap_real_mat (gen_corr_cholesky_unwrapped n) let gen_cov_matrix n = let cov = gen_cov_unwrapped n in wrap_real_mat cov +;; let gen_corr_matrix n = let corr_chol = gen_corr_cholesky_unwrapped n in wrap_real_mat (matprod corr_chol (transpose corr_chol)) +;; let gen_matrix mm m n t = let open Program in @@ -201,12 +231,12 @@ let gen_matrix mm m n t = | CholeskyCov -> gen_cov_cholesky m n | CholeskyCorr -> gen_corr_cholesky m | _ -> - { int_two with - expr= RowVectorExpr (repeat_th m (fun () -> gen_row_vector mm n t)) } + { int_two with expr = RowVectorExpr (repeat_th m (fun () -> gen_row_vector mm n t)) } +;; (* TODO: do some proper random generation of these special matrices *) -let gen_array elt n _ = {int_two with expr= ArrayExpr (repeat_th n elt)} +let gen_array elt n _ = { int_two with expr = ArrayExpr (repeat_th n elt) } let rec generate_value m st t = match st with @@ -214,40 +244,41 @@ let rec generate_value m st t = | SReal -> gen_real m t | SVector e -> gen_vector m (unwrap_int_exn m e) t | SRowVector e -> gen_row_vector m (unwrap_int_exn m e) t - | SMatrix (e1, e2) -> - gen_matrix m (unwrap_int_exn m e1) (unwrap_int_exn m e2) t + | SMatrix (e1, e2) -> gen_matrix m (unwrap_int_exn m e1) (unwrap_int_exn m e2) t | SArray (st, e) -> - let element () = generate_value m st t in - gen_array element (unwrap_int_exn m e) t + let element () = generate_value m st t in + gen_array element (unwrap_int_exn m e) t +;; let rec pp_value_json ppf e = match e.expr with | PostfixOp (e, Transpose) -> pp_value_json ppf e | IntNumeral s | RealNumeral s -> Fmt.string ppf s | ArrayExpr l | RowVectorExpr l -> - Fmt.(pf ppf "[@[%a@]]" (list ~sep:comma pp_value_json) l) + Fmt.(pf ppf "[@[%a@]]" (list ~sep:comma pp_value_json) l) | _ -> failwith "This should never happen." +;; let var_decl_id d = match d.stmt with - | VarDecl {identifier; _} -> identifier.name + | VarDecl { identifier; _ } -> identifier.name | _ -> failwith "This should never happen." +;; let var_decl_gen_val m d = match d.stmt with - | VarDecl {decl_type= Sized sizedtype; transformation; _} -> - generate_value m sizedtype transformation + | VarDecl { decl_type = Sized sizedtype; transformation; _ } -> + generate_value m sizedtype transformation | _ -> failwith "This should never happen." +;; let print_data_prog s = let data = Option.value ~default:[] s.datablock in let l, _ = List.fold data ~init:([], Map.Poly.empty) ~f:(fun (l, m) decl -> let value = var_decl_gen_val m decl in - ( l @ [(var_decl_id decl, value)] - , Map.set m ~key:(var_decl_id decl) ~data:value ) ) - in - let pp ppf (id, value) = - Fmt.pf ppf {|@["%s":@ %a@]|} id pp_value_json value + l @ [ var_decl_id decl, value ], Map.set m ~key:(var_decl_id decl) ~data:value) in + let pp ppf (id, value) = Fmt.pf ppf {|@["%s":@ %a@]|} id pp_value_json value in Fmt.(strf "{@ @[%a@]@ }" (list ~sep:comma pp) l) +;; diff --git a/src/frontend/Debugging.ml b/src/frontend/Debugging.ml index 9c2ef0a67d..3306a6768e 100644 --- a/src/frontend/Debugging.ml +++ b/src/frontend/Debugging.ml @@ -19,9 +19,5 @@ let ast_logger t = if !ast_printing then print_endline (ast_to_string t) (* Controls whether a decorated AST gets printed after the semantic check *) let typed_ast_printing = ref false - -let typed_ast_to_string x = - [%sexp (x : Ast.typed_program)] |> Sexp.to_string_hum - -let typed_ast_logger t = - if !typed_ast_printing then print_endline (typed_ast_to_string t) +let typed_ast_to_string x = [%sexp (x : Ast.typed_program)] |> Sexp.to_string_hum +let typed_ast_logger t = if !typed_ast_printing then print_endline (typed_ast_to_string t) diff --git a/src/frontend/Errors.ml b/src/frontend/Errors.ml index 594de11e4d..2b17d3f84d 100644 --- a/src/frontend/Errors.ml +++ b/src/frontend/Errors.ml @@ -24,52 +24,75 @@ exception FatalError of string (* A fatal error reported by the toplevel *) let fatal_error ?(msg = "") _ = raise (FatalError ("This should never happen. Please file a bug. " ^ msg)) +;; (** Return two lines before and after the specified location and print a message *) let pp_context_and_message ppf (message, loc) = - Fmt.pf ppf "@[%a@,%s@,@]" (Fmt.option Fmt.string) + Fmt.pf + ppf + "@[%a@,%s@,@]" + (Fmt.option Fmt.string) (Location.context_to_string loc) message +;; let pp_semantic_error ppf (message, loc_span) = - Fmt.pf ppf "@[@;Semantic error in %s:@;%a@]@." + Fmt.pf + ppf + "@[@;Semantic error in %s:@;%a@]@." (Location_span.to_string loc_span) pp_context_and_message (message, loc_span.begin_loc) +;; (** A syntax error message used when handling a SyntaxError *) let pp_syntax_error ppf = function | Parsing (message, loc_span) -> - Fmt.pf ppf "@[@,Syntax error in %s, parsing error:@,%a@]@." - (Location_span.to_string loc_span) - pp_context_and_message - (message, loc_span.end_loc) + Fmt.pf + ppf + "@[@,Syntax error in %s, parsing error:@,%a@]@." + (Location_span.to_string loc_span) + pp_context_and_message + (message, loc_span.end_loc) | Lexing (_, loc) -> - Fmt.pf ppf "@[@,Syntax error in %s, lexing error:@,%a@]@." - (Location.to_string {loc with col_num= loc.col_num - 1}) - pp_context_and_message - ("Invalid character found.", loc) + Fmt.pf + ppf + "@[@,Syntax error in %s, lexing error:@,%a@]@." + (Location.to_string { loc with col_num = loc.col_num - 1 }) + pp_context_and_message + ("Invalid character found.", loc) | Include (message, loc) -> - Fmt.pf ppf "@[@,Syntax error in %s, include error:@,%a@]@." - (Location.to_string loc) pp_context_and_message (message, loc) + Fmt.pf + ppf + "@[@,Syntax error in %s, include error:@,%a@]@." + (Location.to_string loc) + pp_context_and_message + (message, loc) +;; (** Switch to control whether warning messages should be printed to stderr (or discarded in case set to false) *) let print_warnings = ref true let without_warnings function_name args = - print_warnings := false ; + print_warnings := false; let out = function_name args in - print_warnings := true ; + print_warnings := true; out +;; (* Warn that a language feature is deprecated *) let warn_deprecated (pos, message) = let loc = - Location.of_position_opt {pos with Lexing.pos_cnum= pos.Lexing.pos_cnum - 1} + Location.of_position_opt { pos with Lexing.pos_cnum = pos.Lexing.pos_cnum - 1 } |> Option.value ~default:Location.empty in - if !print_warnings then - Fmt.pf Fmt.stderr + if !print_warnings + then + Fmt.pf + Fmt.stderr "@[@,Warning: deprecated language construct used in %s:@,%a@]@." - (Location.to_string loc) pp_context_and_message (message, loc) + (Location.to_string loc) + pp_context_and_message + (message, loc) +;; diff --git a/src/frontend/Errors.mli b/src/frontend/Errors.mli index 227f1b1eb8..5fb6f4ae97 100644 --- a/src/frontend/Errors.mli +++ b/src/frontend/Errors.mli @@ -18,17 +18,17 @@ exception SemanticError of (string * Location_span.t) so we can trace their origin. *) exception FatalError of string -val fatal_error : ?msg:string -> unit -> 'a (** Throw a fatal error reported by the toplevel *) +val fatal_error : ?msg:string -> unit -> 'a -val pp_syntax_error : Format.formatter -> syntax_error -> unit (** A syntax error message used when handling a SyntaxError *) +val pp_syntax_error : Format.formatter -> syntax_error -> unit -val pp_semantic_error : Format.formatter -> string * Location_span.t -> unit (** A semantic error message used when handling a SemanticError *) +val pp_semantic_error : Format.formatter -> string * Location_span.t -> unit -val warn_deprecated : Lexing.position * string -> unit (** Warn that a language construct is deprecated *) +val warn_deprecated : Lexing.position * string -> unit -val without_warnings : ('a -> 'b) -> 'a -> 'b (** Evaluate a function on an argument while discarding any warning messages generated by it (rather than printing them to stderr) *) +val without_warnings : ('a -> 'b) -> 'a -> 'b diff --git a/src/frontend/Frontend_utils.ml b/src/frontend/Frontend_utils.ml index 50559f478c..325d79cfc1 100644 --- a/src/frontend/Frontend_utils.ml +++ b/src/frontend/Frontend_utils.ml @@ -3,36 +3,40 @@ open Core_kernel let typed_ast_of_string_exn s = Parse.parse_string Parser.Incremental.program s |> Result.map_error ~f:(Fmt.to_to_string Errors.pp_syntax_error) - |> Result.ok_or_failwith |> Semantic_check.semantic_check_program - |> Result.map_error - ~f:(Fmt.to_to_string (Fmt.list ~sep:Fmt.cut Semantic_error.pp)) |> Result.ok_or_failwith + |> Semantic_check.semantic_check_program + |> Result.map_error ~f:(Fmt.to_to_string (Fmt.list ~sep:Fmt.cut Semantic_error.pp)) + |> Result.ok_or_failwith +;; let get_ast_or_exit filename = try match Parse.parse_file Parser.Incremental.program filename with | Result.Ok ast -> ast | Result.Error err -> - Errors.pp_syntax_error Fmt.stderr err ; - exit 1 - with Errors.SyntaxError err -> - Errors.pp_syntax_error Fmt.stderr err ; + Errors.pp_syntax_error Fmt.stderr err; + exit 1 + with + | Errors.SyntaxError err -> + Errors.pp_syntax_error Fmt.stderr err; exit 1 +;; let type_ast_or_exit ast = try match Semantic_check.semantic_check_program ast with | Result.Ok prog -> prog | Result.Error (error :: _) -> - let loc = Semantic_error.location error - and msg = Fmt.strf "%a" Semantic_error.pp error in - Errors.pp_semantic_error Fmt.stderr (msg, loc) ; - exit 1 + let loc = Semantic_error.location error + and msg = Fmt.strf "%a" Semantic_error.pp error in + Errors.pp_semantic_error Fmt.stderr (msg, loc); + exit 1 | Result.Error [] -> - Printf.eprintf - "Semantic check failed but reported no errors. This should never \ - happen." ; - exit 1 - with Errors.SemanticError err -> - Errors.pp_semantic_error Fmt.stderr err ; + Printf.eprintf + "Semantic check failed but reported no errors. This should never happen."; + exit 1 + with + | Errors.SemanticError err -> + Errors.pp_semantic_error Fmt.stderr err; exit 1 +;; diff --git a/src/frontend/Parse.ml b/src/frontend/Parse.ml index b9b499f754..84e211b70a 100644 --- a/src/frontend/Parse.ml +++ b/src/frontend/Parse.ml @@ -9,9 +9,10 @@ let parse parse_fun lexbuf = error messages support *) let open MenhirLib.General in let module Interp = Parser.MenhirInterpreter in - Stack.push Preprocessor.include_stack lexbuf ; + Stack.push Preprocessor.include_stack lexbuf; let input _ = - (Interp.lexer_lexbuf_to_supplier Lexer.token + (Interp.lexer_lexbuf_to_supplier + Lexer.token (Stack.top_exn Preprocessor.include_stack)) () in @@ -24,63 +25,59 @@ let parse parse_fun lexbuf = in match Interp.stack env with | (lazy Nil) -> - let message = - "Expected \"functions {\" or \"data {\" or \"transformed data {\" \ - or \"parameters {\" or \"transformed parameters {\" or \"model {\" \ - or \"generated quantities {\".\n" - in - Errors.Parsing - ( message - , Option.value_exn - (Location_span.of_positions_opt - (Lexing.lexeme_start_p - (Stack.top_exn Preprocessor.include_stack)) - (Lexing.lexeme_end_p - (Stack.top_exn Preprocessor.include_stack))) ) - |> Result.Error + let message = + "Expected \"functions {\" or \"data {\" or \"transformed data {\" or \ + \"parameters {\" or \"transformed parameters {\" or \"model {\" or \"generated \ + quantities {\".\n" + in + Errors.Parsing + ( message + , Option.value_exn + (Location_span.of_positions_opt + (Lexing.lexeme_start_p (Stack.top_exn Preprocessor.include_stack)) + (Lexing.lexeme_end_p (Stack.top_exn Preprocessor.include_stack))) ) + |> Result.Error | (lazy (Cons (Interp.Element (state, _, start_pos, end_pos), _))) -> - let message = - try - Parsing_errors.message (Interp.number state) - ^ - if !Debugging.grammar_logging then - "(Parse error state " ^ string_of_int (Interp.number state) ^ ")" - else "" - with - | Not_found_s _ -> - if !Debugging.grammar_logging then - "(Parse error state " - ^ string_of_int (Interp.number state) - ^ ")" - else "" - | _ -> - "(Parse error state " ^ string_of_int (Interp.number state) ^ ")" - in - Errors.Parsing - (message, Location_span.of_positions_exn start_pos end_pos) - |> Result.Error + let message = + try + Parsing_errors.message (Interp.number state) + ^ + if !Debugging.grammar_logging + then "(Parse error state " ^ string_of_int (Interp.number state) ^ ")" + else "" + with + | Not_found_s _ -> + if !Debugging.grammar_logging + then "(Parse error state " ^ string_of_int (Interp.number state) ^ ")" + else "" + | _ -> "(Parse error state " ^ string_of_int (Interp.number state) ^ ")" + in + Errors.Parsing (message, Location_span.of_positions_exn start_pos end_pos) + |> Result.Error in Interp.loop_handle success failure input (parse_fun lexbuf.Lexing.lex_curr_p) +;; let parse_string parse_fun str = let lexbuf = let open Lexing in let lexbuf = from_string str in lexbuf.lex_start_p - <- {pos_fname= "string"; pos_lnum= 1; pos_bol= 0; pos_cnum= 0} ; - lexbuf.lex_curr_p <- lexbuf.lex_start_p ; + <- { pos_fname = "string"; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 }; + lexbuf.lex_curr_p <- lexbuf.lex_start_p; lexbuf in parse parse_fun lexbuf +;; let parse_file parse_fun path = let chan = In_channel.create path in let lexbuf = let open Lexing in let lexbuf = from_channel chan in - lexbuf.lex_start_p - <- {pos_fname= path; pos_lnum= 1; pos_bol= 0; pos_cnum= 0} ; - lexbuf.lex_curr_p <- lexbuf.lex_start_p ; + lexbuf.lex_start_p <- { pos_fname = path; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 }; + lexbuf.lex_curr_p <- lexbuf.lex_start_p; lexbuf in parse parse_fun lexbuf +;; diff --git a/src/frontend/Parse.mli b/src/frontend/Parse.mli index 65a33a2b03..435c9f9390 100644 --- a/src/frontend/Parse.mli +++ b/src/frontend/Parse.mli @@ -2,16 +2,16 @@ API *) open Core_kernel -val parse_file : - (Lexing.position -> Ast.untyped_program Parser.MenhirInterpreter.checkpoint) - -> string - -> (Ast.untyped_program, Errors.syntax_error) result (** A helper function to take a parser, a filename and produce an AST. Under the hood, it takes care of Menhir's custom syntax error messages. *) - -val parse_string : - (Lexing.position -> Ast.untyped_program Parser.MenhirInterpreter.checkpoint) +val parse_file + : (Lexing.position -> Ast.untyped_program Parser.MenhirInterpreter.checkpoint) -> string -> (Ast.untyped_program, Errors.syntax_error) result + (** A helper function to take a parser, a string and produce an AST. Under the hood, it takes care of Menhir's custom syntax error messages. *) +val parse_string + : (Lexing.position -> Ast.untyped_program Parser.MenhirInterpreter.checkpoint) + -> string + -> (Ast.untyped_program, Errors.syntax_error) result diff --git a/src/frontend/Preprocessor.ml b/src/frontend/Preprocessor.ml index 734c535c40..50c38e905f 100644 --- a/src/frontend/Preprocessor.ml +++ b/src/frontend/Preprocessor.ml @@ -8,6 +8,7 @@ let dup_exists l = match List.find_a_dup ~compare:String.compare l with | Some _ -> true | None -> false +;; let include_stack = Stack.create () let include_paths : string list ref = ref [] @@ -15,43 +16,46 @@ let include_paths : string list ref = ref [] let rec try_open_in paths fname pos = match paths with | [] -> - raise - (Errors.SyntaxError - (Include - ( "Could not find include file " ^ fname - ^ " in specified include paths.\n" - , Middle.Location.of_position_exn - (lexeme_start_p (Stack.top_exn include_stack)) ))) - | path :: rest_of_paths -> ( - try - let full_path = path ^ "/" ^ fname in - ( In_channel.create full_path - , sprintf "%s, included from\n%s" full_path - (Middle.Location.to_string - (Middle.Location.of_position_exn - (Stack.top_exn include_stack).lex_start_p)) ) - with _ -> try_open_in rest_of_paths fname pos ) + raise + (Errors.SyntaxError + (Include + ( "Could not find include file " ^ fname ^ " in specified include paths.\n" + , Middle.Location.of_position_exn + (lexeme_start_p (Stack.top_exn include_stack)) ))) + | path :: rest_of_paths -> + (try + let full_path = path ^ "/" ^ fname in + ( In_channel.create full_path + , sprintf + "%s, included from\n%s" + full_path + (Middle.Location.to_string + (Middle.Location.of_position_exn (Stack.top_exn include_stack).lex_start_p)) + ) + with + | _ -> try_open_in rest_of_paths fname pos) +;; let maybe_remove_quotes str = let open String in - if is_prefix str ~prefix:"\"" && is_suffix str ~suffix:"\"" then - drop_suffix (drop_prefix str 1) 1 + if is_prefix str ~prefix:"\"" && is_suffix str ~suffix:"\"" + then drop_suffix (drop_prefix str 1) 1 else str +;; let try_get_new_lexbuf fname pos = - let chan, path = - try_open_in !include_paths (maybe_remove_quotes fname) pos - in + let chan, path = try_open_in !include_paths (maybe_remove_quotes fname) pos in let new_lexbuf = from_channel chan in - new_lexbuf.lex_start_p - <- {pos_fname= path; pos_lnum= 1; pos_bol= 0; pos_cnum= 0} ; - new_lexbuf.lex_curr_p <- new_lexbuf.lex_start_p ; - if dup_exists (Str.split (Str.regexp ", included from\n") path) then + new_lexbuf.lex_start_p <- { pos_fname = path; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 }; + new_lexbuf.lex_curr_p <- new_lexbuf.lex_start_p; + if dup_exists (Str.split (Str.regexp ", included from\n") path) + then raise (Errors.SyntaxError (Include ( Printf.sprintf "File %s recursively included itself.\n" fname , Middle.Location.of_position_exn - (lexeme_start_p (Stack.top_exn include_stack)) ))) ; - Stack.push include_stack new_lexbuf ; + (lexeme_start_p (Stack.top_exn include_stack)) ))); + Stack.push include_stack new_lexbuf; new_lexbuf +;; diff --git a/src/frontend/Preprocessor.mli b/src/frontend/Preprocessor.mli index 4ba7accd8c..d5195b6490 100644 --- a/src/frontend/Preprocessor.mli +++ b/src/frontend/Preprocessor.mli @@ -2,13 +2,13 @@ open Core_kernel -val include_stack : Lexing.lexbuf Stack.t (** Stack with lexing buffers, created from all the includes encountered so far *) +val include_stack : Lexing.lexbuf Stack.t -val include_paths : string list ref (** List of paths to search for including files *) +val include_paths : string list ref -val try_get_new_lexbuf : string -> Lexing.position -> Lexing.lexbuf (** Search include paths for filename and try to create a new lexing buffer with that filename, record that included from specified position *) +val try_get_new_lexbuf : string -> Lexing.position -> Lexing.lexbuf diff --git a/src/frontend/Pretty_printing.ml b/src/frontend/Pretty_printing.ml index 50efb80f24..7c2dea89e1 100644 --- a/src/frontend/Pretty_printing.ml +++ b/src/frontend/Pretty_printing.ml @@ -10,21 +10,28 @@ let wrap_fmt fmt x = https://discuss.ocaml.org/t/debugging-memory-issues/3223/8 *) Fmt.strf "%a" fmt x +;; let with_hbox ppf f = - Format.pp_open_hbox ppf () ; f () ; Format.pp_close_box ppf () ; () + Format.pp_open_hbox ppf (); + f (); + Format.pp_close_box ppf (); + () +;; let with_box ppf offset f = - Format.pp_open_box ppf offset ; - f () ; - Format.pp_close_box ppf () ; + Format.pp_open_box ppf offset; + f (); + Format.pp_close_box ppf (); () +;; let with_vbox ppf offset f = - Format.pp_open_vbox ppf offset ; - f () ; - Format.pp_close_box ppf () ; + Format.pp_open_vbox ppf offset; + f (); + Format.pp_close_box ppf (); () +;; let comma_no_break = Fmt.unit ", " @@ -32,23 +39,28 @@ let with_indented_box ppf indentation offset f = let rec pp_print_n_spaces ppf = function | 0 -> () | i -> - Format.pp_print_space ppf () ; - pp_print_n_spaces ppf (i - 1) + Format.pp_print_space ppf (); + pp_print_n_spaces ppf (i - 1) in with_hbox ppf (fun () -> - pp_print_n_spaces ppf indentation ; - with_box ppf offset f ) ; + pp_print_n_spaces ppf indentation; + with_box ppf offset f); () +;; let rec unwind_sized_array_type = function - | Middle.SizedType.SArray (st, e) -> ( - match unwind_sized_array_type st with st2, es -> (st2, es @ [e]) ) - | st -> (st, []) + | Middle.SizedType.SArray (st, e) -> + (match unwind_sized_array_type st with + | st2, es -> st2, es @ [ e ]) + | st -> st, [] +;; let rec unwind_array_type = function - | Middle.UnsizedType.UArray ut -> ( - match unwind_array_type ut with ut2, d -> (ut2, d + 1) ) - | ut -> (ut, 0) + | Middle.UnsizedType.UArray ut -> + (match unwind_array_type ut with + | ut2, d -> ut2, d + 1) + | ut -> ut, 0 +;; (** XXX this should use the MIR pretty printers after AST pretty printers are updated to use `Fmt`. *) @@ -63,13 +75,17 @@ and pp_unsizedtype ppf = function | URowVector -> Fmt.pf ppf "row_vector" | UMatrix -> Fmt.pf ppf "matrix" | UArray ut -> - let ut2, d = unwind_array_type ut in - let array_str = "[" ^ String.make d ',' ^ "]" in - Fmt.(suffix (const string array_str) pp_unsizedtype ppf ut2) + let ut2, d = unwind_array_type ut in + let array_str = "[" ^ String.make d ',' ^ "]" in + Fmt.(suffix (const string array_str) pp_unsizedtype ppf ut2) | UFun (argtypes, rt) -> - Fmt.pf ppf "{|@[(%a) => %a@]|}" - Fmt.(list ~sep:comma_no_break pp_argtype) - argtypes pp_returntype rt + Fmt.pf + ppf + "{|@[(%a) => %a@]|}" + Fmt.(list ~sep:comma_no_break pp_argtype) + argtypes + pp_returntype + rt | UMathLibraryFunction -> Fmt.pf ppf "Stan Math function" and pp_unsizedtypes ppf l = Fmt.(list ~sep:comma_no_break pp_unsizedtype) ppf l @@ -113,52 +129,59 @@ and pp_index ppf = function | Between (e1, e2) -> Fmt.pf ppf "%a : %a" pp_expression e1 pp_expression e2 and pp_list_of_indices ppf l = - Fmt.(list ~sep:comma_no_break pp_index) ppf l ; + Fmt.(list ~sep:comma_no_break pp_index) ppf l; () -and pp_expression ppf {expr= e_content; _} = +and pp_expression ppf { expr = e_content; _ } = match e_content with | TernaryIf (e1, e2, e3) -> - with_box ppf 0 (fun () -> - Fmt.pf ppf "%a" pp_expression e1 ; - Format.pp_print_space ppf () ; - Fmt.pf ppf "? %a" pp_expression e2 ; - Format.pp_print_space ppf () ; - Fmt.pf ppf ": %a" pp_expression e3 ) + with_box ppf 0 (fun () -> + Fmt.pf ppf "%a" pp_expression e1; + Format.pp_print_space ppf (); + Fmt.pf ppf "? %a" pp_expression e2; + Format.pp_print_space ppf (); + Fmt.pf ppf ": %a" pp_expression e3) | BinOp (e1, op, e2) -> - with_box ppf 0 (fun () -> - Fmt.pf ppf "%a" pp_expression e1 ; - Format.pp_print_space ppf () ; - Fmt.pf ppf "%a %a" pp_operator op pp_expression e2 ) + with_box ppf 0 (fun () -> + Fmt.pf ppf "%a" pp_expression e1; + Format.pp_print_space ppf (); + Fmt.pf ppf "%a %a" pp_operator op pp_expression e2) | PrefixOp (op, e) -> Fmt.pf ppf "%a%a" pp_operator op pp_expression e | PostfixOp (e, op) -> Fmt.pf ppf "%a%a" pp_expression e pp_operator op | Variable id -> pp_identifier ppf id | IntNumeral i -> Fmt.pf ppf "%s" i | RealNumeral r -> Fmt.pf ppf "%s" r | FunApp (_, id, es) -> - Fmt.pf ppf "%a(" pp_identifier id ; - with_box ppf 0 (fun () -> Fmt.pf ppf "%a)" pp_list_of_expression es) - | CondDistApp (_, id, es) -> ( - match es with + Fmt.pf ppf "%a(" pp_identifier id; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a)" pp_list_of_expression es) + | CondDistApp (_, id, es) -> + (match es with | [] -> Errors.fatal_error () | e :: es' -> - with_hbox ppf (fun () -> - Fmt.pf ppf "%a(%a| %a)" pp_identifier id pp_expression e - pp_list_of_expression es' ) ) + with_hbox ppf (fun () -> + Fmt.pf + ppf + "%a(%a| %a)" + pp_identifier + id + pp_expression + e + pp_list_of_expression + es')) (* GetLP is deprecated *) | GetLP -> Fmt.pf ppf "get_lp()" | GetTarget -> Fmt.pf ppf "target()" | ArrayExpr es -> - Fmt.pf ppf "{" ; - with_box ppf 0 (fun () -> Fmt.pf ppf "%a}" pp_list_of_expression es) + Fmt.pf ppf "{"; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a}" pp_list_of_expression es) | RowVectorExpr es -> - Fmt.pf ppf "[" ; - with_box ppf 0 (fun () -> Fmt.pf ppf "%a]" pp_list_of_expression es) + Fmt.pf ppf "["; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a]" pp_list_of_expression es) | Paren e -> Fmt.pf ppf "(%a)" pp_expression e - | Indexed (e, l) -> ( - match l with + | Indexed (e, l) -> + (match l with | [] -> Fmt.pf ppf "%a" pp_expression e - | l -> Fmt.pf ppf "%a[%a]" pp_expression e pp_list_of_indices l ) + | l -> Fmt.pf ppf "%a[%a]" pp_expression e pp_list_of_indices l) and pp_list_of_expression ppf es = Fmt.(list ~sep:comma pp_expression) ppf es and pp_lvalue ppf lhs = pp_expression ppf (expr_of_lvalue lhs) @@ -175,23 +198,20 @@ and pp_truncation ppf = function | NoTruncate -> Fmt.pf ppf "" | TruncateUpFrom e -> Fmt.pf ppf " T[%a, ]" pp_expression e | TruncateDownFrom e -> Fmt.pf ppf " T[ , %a]" pp_expression e - | TruncateBetween (e1, e2) -> - Fmt.pf ppf " T[%a, %a]" pp_expression e1 pp_expression e2 + | TruncateBetween (e1, e2) -> Fmt.pf ppf " T[%a, %a]" pp_expression e1 pp_expression e2 and pp_printable ppf = function | PString s -> Fmt.pf ppf "%s" s | PExpr e -> pp_expression ppf e -and pp_list_of_printables ppf l = - Fmt.(list ~sep:comma_no_break pp_printable) ppf l +and pp_list_of_printables ppf l = Fmt.(list ~sep:comma_no_break pp_printable) ppf l and pp_sizedtype ppf = function | Middle.SizedType.SInt -> Fmt.pf ppf "int" | SReal -> Fmt.pf ppf "real" | SVector e -> Fmt.pf ppf "vector[%a]" pp_expression e | SRowVector e -> Fmt.pf ppf "row_vector[%a]" pp_expression e - | SMatrix (e1, e2) -> - Fmt.pf ppf "matrix[%a, %a]" pp_expression e1 pp_expression e2 + | SMatrix (e1, e2) -> Fmt.pf ppf "matrix[%a, %a]" pp_expression e1 pp_expression e2 | SArray _ -> raise (Errors.FatalError "This should never happen.") and pp_transformation ppf = function @@ -199,11 +219,11 @@ and pp_transformation ppf = function | Lower e -> Fmt.pf ppf "" pp_expression e | Upper e -> Fmt.pf ppf "" pp_expression e | LowerUpper (e1, e2) -> - Fmt.pf ppf "" pp_expression e1 pp_expression e2 + Fmt.pf ppf "" pp_expression e1 pp_expression e2 | Offset e -> Fmt.pf ppf "" pp_expression e | Multiplier e -> Fmt.pf ppf "" pp_expression e | OffsetMultiplier (e1, e2) -> - Fmt.pf ppf "" pp_expression e1 pp_expression e2 + Fmt.pf ppf "" pp_expression e1 pp_expression e2 | Ordered -> Fmt.pf ppf "" | PositiveOrdered -> Fmt.pf ppf "" | Simplex -> Fmt.pf ppf "" @@ -217,7 +237,7 @@ and pp_transformed_type ppf (pst, trans) = let rec discard_arrays pst = match pst with | Middle.Type.Sized st -> - Middle.Type.Sized (Fn.compose fst unwind_sized_array_type st) + Middle.Type.Sized (Fn.compose fst unwind_sized_array_type st) | Unsized (UArray t) -> discard_arrays (Unsized t) | Unsized ut -> Unsized ut in @@ -225,39 +245,29 @@ and pp_transformed_type ppf (pst, trans) = let unsizedtype_fmt = match pst with | Middle.Type.Sized (SArray _ as st) -> - Fmt.const pp_sizedtype (Fn.compose fst unwind_sized_array_type st) + Fmt.const pp_sizedtype (Fn.compose fst unwind_sized_array_type st) | _ -> Fmt.const pp_unsizedtype (Middle.Type.to_unsized pst) in let sizes_fmt = match pst with | Sized (SVector e) | Sized (SRowVector e) -> - Fmt.const (fun ppf -> Fmt.pf ppf "[%a]" pp_expression) e + Fmt.const (fun ppf -> Fmt.pf ppf "[%a]" pp_expression) e | Sized (SMatrix (e1, e2)) -> - Fmt.const - (fun ppf -> Fmt.pf ppf "[%a, %a]" pp_expression e1 pp_expression) - e2 - | Sized (SArray _) | Unsized _ | Sized Middle.SizedType.SInt | Sized SReal - -> - Fmt.nop + Fmt.const (fun ppf -> Fmt.pf ppf "[%a, %a]" pp_expression e1 pp_expression) e2 + | Sized (SArray _) | Unsized _ | Sized Middle.SizedType.SInt | Sized SReal -> Fmt.nop in let cov_sizes_fmt = match pst with | Sized (SMatrix (e1, e2)) -> - if e1 = e2 then - Fmt.const (fun ppf -> Fmt.pf ppf "[%a]" pp_expression) e1 - else - Fmt.const - (fun ppf -> Fmt.pf ppf "[%a, %a]" pp_expression e1 pp_expression) - e2 + if e1 = e2 + then Fmt.const (fun ppf -> Fmt.pf ppf "[%a]" pp_expression) e1 + else Fmt.const (fun ppf -> Fmt.pf ppf "[%a, %a]" pp_expression e1 pp_expression) e2 | _ -> Fmt.nop in match trans with - | Middle.Program.Identity -> - Fmt.pf ppf "%a%a" unsizedtype_fmt () sizes_fmt () - | Lower _ | Upper _ | LowerUpper _ | Offset _ | Multiplier _ - |OffsetMultiplier _ -> - Fmt.pf ppf "%a%a%a" unsizedtype_fmt () pp_transformation trans sizes_fmt - () + | Middle.Program.Identity -> Fmt.pf ppf "%a%a" unsizedtype_fmt () sizes_fmt () + | Lower _ | Upper _ | LowerUpper _ | Offset _ | Multiplier _ | OffsetMultiplier _ -> + Fmt.pf ppf "%a%a%a" unsizedtype_fmt () pp_transformation trans sizes_fmt () | Ordered -> Fmt.pf ppf "ordered%a" sizes_fmt () | PositiveOrdered -> Fmt.pf ppf "positive_ordered%a" sizes_fmt () | Simplex -> Fmt.pf ppf "simplex%a" sizes_fmt () @@ -270,16 +280,15 @@ and pp_transformed_type ppf (pst, trans) = and pp_array_dims ppf = function | [] -> Fmt.pf ppf "" | es -> - Fmt.pf ppf "[" ; - with_box ppf 0 (fun () -> - Fmt.pf ppf "%a]" pp_list_of_expression (List.rev es) ) + Fmt.pf ppf "["; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a]" pp_list_of_expression (List.rev es)) and pp_indent_unless_block ppf s = match s.stmt with | Block _ -> pp_statement ppf s | _ -> - Format.pp_print_cut ppf () ; - with_indented_box ppf 2 0 (fun () -> Fmt.pf ppf "%a" pp_statement s) + Format.pp_print_cut ppf (); + with_indented_box ppf 2 0 (fun () -> Fmt.pf ppf "%a" pp_statement s) (* This function helps write chained if-then-else-if-... blocks * correctly. Without it, each IfThenElse would trigger a new @@ -288,135 +297,165 @@ and pp_indent_unless_block ppf s = and pp_recursive_ifthenelse ppf s = match s.stmt with | IfThenElse (e, s, None) -> - Fmt.pf ppf "if (%a) %a" pp_expression e pp_indent_unless_block s + Fmt.pf ppf "if (%a) %a" pp_expression e pp_indent_unless_block s | IfThenElse (e, s1, Some s2) -> - Fmt.pf ppf "if (%a) %a" pp_expression e pp_indent_unless_block s1 ; - Format.pp_print_cut ppf () ; - Fmt.pf ppf "else %a" pp_recursive_ifthenelse s2 + Fmt.pf ppf "if (%a) %a" pp_expression e pp_indent_unless_block s1; + Format.pp_print_cut ppf (); + Fmt.pf ppf "else %a" pp_recursive_ifthenelse s2 | _ -> pp_indent_unless_block ppf s -and pp_statement ppf ({stmt= s_content; _} as ss) = +and pp_statement ppf ({ stmt = s_content; _ } as ss) = match s_content with - | Assignment {assign_lhs= l; assign_op= assop; assign_rhs= e} -> - with_hbox ppf (fun () -> - Fmt.pf ppf "%a %a %a;" pp_lvalue l pp_assignmentoperator assop - pp_expression e ) + | Assignment { assign_lhs = l; assign_op = assop; assign_rhs = e } -> + with_hbox ppf (fun () -> + Fmt.pf ppf "%a %a %a;" pp_lvalue l pp_assignmentoperator assop pp_expression e) | NRFunApp (_, id, es) -> - Fmt.pf ppf "%a(" pp_identifier id ; - with_box ppf 0 (fun () -> Fmt.pf ppf "%a);" pp_list_of_expression es) + Fmt.pf ppf "%a(" pp_identifier id; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a);" pp_list_of_expression es) | TargetPE e -> Fmt.pf ppf "target += %a;" pp_expression e | IncrementLogProb e -> - with_hbox ppf (fun () -> - Fmt.pf ppf "increment_log_prob(%a);" pp_expression e ) - | Tilde {arg= e; distribution= id; args= es; truncation= t} -> - Fmt.pf ppf "%a ~ %a(" pp_expression e pp_identifier id ; - with_box ppf 0 (fun () -> Fmt.pf ppf "%a)" pp_list_of_expression es) ; - Fmt.pf ppf "%a;" pp_truncation t + with_hbox ppf (fun () -> Fmt.pf ppf "increment_log_prob(%a);" pp_expression e) + | Tilde { arg = e; distribution = id; args = es; truncation = t } -> + Fmt.pf ppf "%a ~ %a(" pp_expression e pp_identifier id; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a)" pp_list_of_expression es); + Fmt.pf ppf "%a;" pp_truncation t | Break -> Fmt.pf ppf "break;" | Continue -> Fmt.pf ppf "continue;" - | Return e -> - with_hbox ppf (fun () -> Fmt.pf ppf "return %a;" pp_expression e) + | Return e -> with_hbox ppf (fun () -> Fmt.pf ppf "return %a;" pp_expression e) | ReturnVoid -> Fmt.pf ppf "return;" | Print ps -> Fmt.pf ppf "print(%a);" pp_list_of_printables ps | Reject ps -> Fmt.pf ppf "reject(%a);" pp_list_of_printables ps | Skip -> Fmt.pf ppf ";" - | IfThenElse (_, _, _) -> - with_vbox ppf 0 (fun () -> pp_recursive_ifthenelse ppf ss) + | IfThenElse (_, _, _) -> with_vbox ppf 0 (fun () -> pp_recursive_ifthenelse ppf ss) | While (e, s) -> Fmt.pf ppf "while (%a) %a" pp_expression e pp_statement s - | For {loop_variable= id; lower_bound= e1; upper_bound= e2; loop_body= s} -> - with_vbox ppf 0 (fun () -> - Fmt.pf ppf "for (%a in %a : %a) %a" pp_identifier id pp_expression e1 - pp_expression e2 pp_indent_unless_block s ) + | For { loop_variable = id; lower_bound = e1; upper_bound = e2; loop_body = s } -> + with_vbox ppf 0 (fun () -> + Fmt.pf + ppf + "for (%a in %a : %a) %a" + pp_identifier + id + pp_expression + e1 + pp_expression + e2 + pp_indent_unless_block + s) | ForEach (id, e, s) -> - Fmt.pf ppf "for (%a in %a) %a" pp_identifier id pp_expression e - pp_indent_unless_block s + Fmt.pf + ppf + "for (%a in %a) %a" + pp_identifier + id + pp_expression + e + pp_indent_unless_block + s | Block vdsl -> - Fmt.pf ppf "{" ; - Format.pp_print_cut ppf () ; - with_indented_box ppf 2 0 (fun () -> pp_list_of_statements ppf vdsl) ; - Format.pp_print_cut ppf () ; - Fmt.pf ppf "}" + Fmt.pf ppf "{"; + Format.pp_print_cut ppf (); + with_indented_box ppf 2 0 (fun () -> pp_list_of_statements ppf vdsl); + Format.pp_print_cut ppf (); + Fmt.pf ppf "}" | VarDecl - { decl_type= pst - ; transformation= trans - ; identifier= id - ; initial_value= init - ; is_global= _ } -> - let pp_init ppf init = - match init with - | None -> Fmt.pf ppf "" - | Some e -> Fmt.pf ppf " = %a" pp_expression e - in - let es = - match pst with - | Sized st -> Fn.compose snd unwind_sized_array_type st - | Unsized _ -> [] - in - with_hbox ppf (fun () -> - Fmt.pf ppf "%a %a%a%a;" pp_transformed_type (pst, trans) - pp_identifier id pp_array_dims es pp_init init ) - | FunDef {returntype= rt; funname= id; arguments= args; body= b} -> ( - Fmt.pf ppf "%a %a(" pp_returntype rt pp_identifier id ; - with_box ppf 0 (fun () -> - Fmt.pf ppf "%a" (Fmt.list ~sep:Fmt.comma pp_args) args ) ; - match b with - | {stmt= Skip; _} -> Fmt.pf ppf ");" - | b -> Fmt.pf ppf ") %a" pp_statement b ) + { decl_type = pst + ; transformation = trans + ; identifier = id + ; initial_value = init + ; is_global = _ + } -> + let pp_init ppf init = + match init with + | None -> Fmt.pf ppf "" + | Some e -> Fmt.pf ppf " = %a" pp_expression e + in + let es = + match pst with + | Sized st -> Fn.compose snd unwind_sized_array_type st + | Unsized _ -> [] + in + with_hbox ppf (fun () -> + Fmt.pf + ppf + "%a %a%a%a;" + pp_transformed_type + (pst, trans) + pp_identifier + id + pp_array_dims + es + pp_init + init) + | FunDef { returntype = rt; funname = id; arguments = args; body = b } -> + Fmt.pf ppf "%a %a(" pp_returntype rt pp_identifier id; + with_box ppf 0 (fun () -> Fmt.pf ppf "%a" (Fmt.list ~sep:Fmt.comma pp_args) args); + (match b with + | { stmt = Skip; _ } -> Fmt.pf ppf ");" + | b -> Fmt.pf ppf ") %a" pp_statement b) and pp_args ppf (at, ut, id) = Fmt.pf ppf "%a%a %a" pp_autodifftype at pp_unsizedtype ut pp_identifier id and pp_list_of_statements ppf l = with_vbox ppf 0 (fun () -> Format.pp_print_list pp_statement ppf l) +;; let pp_block block_name ppf block_stmts = - Fmt.pf ppf "%s {" block_name ; - Format.pp_print_cut ppf () ; - if List.length block_stmts > 0 then ( + Fmt.pf ppf "%s {" block_name; + Format.pp_print_cut ppf (); + if List.length block_stmts > 0 + then ( with_indented_box ppf 2 0 (fun () -> - pp_list_of_statements ppf block_stmts ; - () ) ; - Format.pp_print_cut ppf () ) - else Format.pp_print_cut ppf () ; - Fmt.pf ppf "}" ; + pp_list_of_statements ppf block_stmts; + ()); + Format.pp_print_cut ppf ()) + else Format.pp_print_cut ppf (); + Fmt.pf ppf "}"; Format.pp_print_cut ppf () +;; let pp_opt_block ppf block_name opt_block = Fmt.option ~none:Fmt.nop (pp_block block_name) ppf opt_block - -let pp_program ppf - { functionblock= bf - ; datablock= bd - ; transformeddatablock= btd - ; parametersblock= bp - ; transformedparametersblock= btp - ; modelblock= bm - ; generatedquantitiesblock= bgq } = - Format.pp_open_vbox ppf 0 ; - pp_opt_block ppf "functions" bf ; - pp_opt_block ppf "data" bd ; - pp_opt_block ppf "transformed data" btd ; - pp_opt_block ppf "parameters" bp ; - pp_opt_block ppf "transformed parameters" btp ; - pp_opt_block ppf "model" bm ; - pp_opt_block ppf "generated quantities" bgq ; +;; + +let pp_program + ppf + { functionblock = bf + ; datablock = bd + ; transformeddatablock = btd + ; parametersblock = bp + ; transformedparametersblock = btp + ; modelblock = bm + ; generatedquantitiesblock = bgq + } + = + Format.pp_open_vbox ppf 0; + pp_opt_block ppf "functions" bf; + pp_opt_block ppf "data" bd; + pp_opt_block ppf "transformed data" btd; + pp_opt_block ppf "parameters" bp; + pp_opt_block ppf "transformed parameters" btp; + pp_opt_block ppf "model" bm; + pp_opt_block ppf "generated quantities" bgq; Format.pp_close_box ppf () +;; let check_correctness prog pretty = let result_ast = - Errors.without_warnings - (Parse.parse_string Parser.Incremental.program) - pretty + Errors.without_warnings (Parse.parse_string Parser.Incremental.program) pretty in - if - compare_untyped_program prog (Option.value_exn (Result.ok result_ast)) <> 0 + if compare_untyped_program prog (Option.value_exn (Result.ok result_ast)) <> 0 then failwith "Pretty printing failed. Please file a bug." +;; let pretty_print_program p = let result = wrap_fmt pp_program p in - check_correctness p result ; result + check_correctness p result; + result +;; let pretty_print_typed_program p = let result = wrap_fmt pp_program p in - check_correctness (untyped_program_of_typed_program p) result ; + check_correctness (untyped_program_of_typed_program p) result; result +;; diff --git a/src/frontend/Semantic_check.ml b/src/frontend/Semantic_check.ml index 94ba415a70..9f7faec72a 100644 --- a/src/frontend/Semantic_check.ml +++ b/src/frontend/Semantic_check.ml @@ -19,16 +19,16 @@ module Validate = Common.Validation.Make (Semantic_error) let check_of_compatible_return_type rt1 srt2 = UnsizedType.( - match (rt1, srt2) with + match rt1, srt2 with | Void, NoReturnType - |Void, Incomplete Void - |Void, Complete Void - |Void, AnyReturnType -> - true + | Void, Incomplete Void + | Void, Complete Void + | Void, AnyReturnType -> true | ReturnType UReal, Complete (ReturnType UInt) -> true | ReturnType rt1, Complete (ReturnType rt2) -> rt1 = rt2 | ReturnType _, AnyReturnType -> true | _ -> false) +;; (** Origin blocks, to keep track of where variables are declared *) type originblock = @@ -50,19 +50,21 @@ let vm = Symbol_table.initialize () (* Record structure holding flags and other markers about context to be used for error reporting. *) type context_flags_record = - { current_block: originblock - ; in_toplevel_decl: bool - ; in_fun_def: bool - ; in_returning_fun_def: bool - ; in_rng_fun_def: bool - ; in_lp_fun_def: bool - ; loop_depth: int } + { current_block : originblock + ; in_toplevel_decl : bool + ; in_fun_def : bool + ; in_returning_fun_def : bool + ; in_rng_fun_def : bool + ; in_lp_fun_def : bool + ; loop_depth : int + } (* Some helper functions *) let dup_exists l = match List.find_a_dup ~compare:String.compare l with | Some _ -> true | None -> false +;; let type_of_expr_typed ue = ue.emeta.type_ @@ -70,44 +72,50 @@ let calculate_autodifftype cf at ut = match at with | (Param | TParam | Model | Functions) when not (UnsizedType.contains_int ut || cf.current_block = GQuant) -> - UnsizedType.AutoDiffable + UnsizedType.AutoDiffable | _ -> DataOnly +;; let has_int_type ue = ue.emeta.type_ = UInt let has_int_array_type ue = ue.emeta.type_ = UArray UInt let has_int_or_real_type ue = - match ue.emeta.type_ with UInt | UReal -> true | _ -> false + match ue.emeta.type_ with + | UInt | UReal -> true + | _ -> false +;; let probability_distribution_name_variants id = let name = id.name in let open String in List.map - ~f:(fun n -> {name= n; id_loc= id.id_loc}) - ( if name = "multiply_log" || name = "binomial_coefficient_log" then [name] - else if is_suffix ~suffix:"_lpmf" name then - [name; drop_suffix name 5 ^ "_lpdf"; drop_suffix name 5 ^ "_log"] - else if is_suffix ~suffix:"_lpdf" name then - [name; drop_suffix name 5 ^ "_lpmf"; drop_suffix name 5 ^ "_log"] - else if is_suffix ~suffix:"_lcdf" name then - [name; drop_suffix name 5 ^ "_cdf_log"] - else if is_suffix ~suffix:"_lccdf" name then - [name; drop_suffix name 6 ^ "_ccdf_log"] - else if is_suffix ~suffix:"_cdf_log" name then - [name; drop_suffix name 8 ^ "_lcdf"] - else if is_suffix ~suffix:"_ccdf_log" name then - [name; drop_suffix name 9 ^ "_lccdf"] - else if is_suffix ~suffix:"_log" name then - [name; drop_suffix name 4 ^ "_lpmf"; drop_suffix name 4 ^ "_lpdf"] - else [name] ) + ~f:(fun n -> { name = n; id_loc = id.id_loc }) + (if name = "multiply_log" || name = "binomial_coefficient_log" + then [ name ] + else if is_suffix ~suffix:"_lpmf" name + then [ name; drop_suffix name 5 ^ "_lpdf"; drop_suffix name 5 ^ "_log" ] + else if is_suffix ~suffix:"_lpdf" name + then [ name; drop_suffix name 5 ^ "_lpmf"; drop_suffix name 5 ^ "_log" ] + else if is_suffix ~suffix:"_lcdf" name + then [ name; drop_suffix name 5 ^ "_cdf_log" ] + else if is_suffix ~suffix:"_lccdf" name + then [ name; drop_suffix name 6 ^ "_ccdf_log" ] + else if is_suffix ~suffix:"_cdf_log" name + then [ name; drop_suffix name 8 ^ "_lcdf" ] + else if is_suffix ~suffix:"_ccdf_log" name + then [ name; drop_suffix name 9 ^ "_lccdf" ] + else if is_suffix ~suffix:"_log" name + then [ name; drop_suffix name 4 ^ "_lpmf"; drop_suffix name 4 ^ "_lpdf" ] + else [ name ]) +;; let lub_rt loc rt1 rt2 = - match (rt1, rt2) with + match rt1, rt2 with | UnsizedType.ReturnType UReal, UnsizedType.ReturnType UInt - |ReturnType UInt, ReturnType UReal -> - Validate.ok (UnsizedType.ReturnType UReal) + | ReturnType UInt, ReturnType UReal -> Validate.ok (UnsizedType.ReturnType UReal) | _, _ when rt1 = rt2 -> Validate.ok rt2 | _ -> Semantic_error.mismatched_return_types loc rt1 rt2 |> Validate.error +;; let check_fresh_variable_basic id is_nullary_function = Validate.( @@ -116,23 +124,24 @@ let check_fresh_variable_basic id is_nullary_function = not of nullary function types to clash with nullary library functions. No other name clashes are tolerated. Here's the logic to achieve that. *) - if - Stan_math_signatures.is_stan_math_function_name id.name - && ( is_nullary_function - || Stan_math_signatures.stan_math_returntype id.name [] = None ) - || Stan_math_signatures.is_reduce_sum_fn id.name + if (Stan_math_signatures.is_stan_math_function_name id.name + && (is_nullary_function + || Stan_math_signatures.stan_math_returntype id.name [] = None)) + || Stan_math_signatures.is_reduce_sum_fn id.name then Semantic_error.ident_is_stanmath_name id.id_loc id.name |> error - else + else ( match Symbol_table.look vm id.name with | Some _ -> Semantic_error.ident_in_use id.id_loc id.name |> error - | None -> ok ()) + | None -> ok ())) +;; let check_fresh_variable id is_nullary_function = - List.fold ~init:(Validate.ok ()) + List.fold + ~init:(Validate.ok ()) ~f:(fun v0 name -> - check_fresh_variable_basic name is_nullary_function - |> Validate.apply_const v0 ) + check_fresh_variable_basic name is_nullary_function |> Validate.apply_const v0) (probability_distribution_name_variants id) +;; (* == SEMANTIC CHECK OF PROGRAM ELEMENTS ==================================== *) @@ -143,156 +152,244 @@ let semantic_check_assignmentoperator op = Validate.ok op let semantic_check_autodifftype at = Validate.ok at (* Probably nothing to do here *) -let rec semantic_check_unsizedtype : UnsizedType.t -> unit Validate.t = - function +let rec semantic_check_unsizedtype : UnsizedType.t -> unit Validate.t = function | UFun (l, rt) -> - (* fold over argument types accumulating errors with initial state + (* fold over argument types accumulating errors with initial state given by validating the return type *) - List.fold - ~f:(fun v0 (at, ut) -> - Validate.( - apply_const - (apply_const v0 (semantic_check_autodifftype at)) - (semantic_check_unsizedtype ut)) ) - ~init:(semantic_check_returntype rt) - l + List.fold + ~f:(fun v0 (at, ut) -> + Validate.( + apply_const + (apply_const v0 (semantic_check_autodifftype at)) + (semantic_check_unsizedtype ut))) + ~init:(semantic_check_returntype rt) + l | UArray ut -> semantic_check_unsizedtype ut | _ -> Validate.ok () -and semantic_check_returntype : UnsizedType.returntype -> unit Validate.t = - function +and semantic_check_returntype : UnsizedType.returntype -> unit Validate.t = function | Void -> Validate.ok () | ReturnType ut -> semantic_check_unsizedtype ut +;; (* -- Indentifiers ---------------------------------------------------------- *) let reserved_keywords = - [ "true"; "false"; "repeat"; "until"; "then"; "var"; "fvar"; "STAN_MAJOR" - ; "STAN_MINOR"; "STAN_PATCH"; "STAN_MATH_MAJOR"; "STAN_MATH_MINOR" - ; "STAN_MATH_PATCH"; "alignas"; "alignof"; "and"; "and_eq"; "asm"; "auto" - ; "bitand"; "bitor"; "bool"; "break"; "case"; "catch"; "char"; "char16_t" - ; "char32_t"; "class"; "compl"; "const"; "constexpr"; "const_cast" - ; "continue"; "decltype"; "default"; "delete"; "do"; "double"; "dynamic_cast" - ; "else"; "enum"; "explicit"; "export"; "extern"; "false"; "float"; "for" - ; "friend"; "goto"; "if"; "inline"; "int"; "long"; "mutable"; "namespace" - ; "new"; "noexcept"; "not"; "not_eq"; "nullptr"; "operator"; "or"; "or_eq" - ; "private"; "protected"; "public"; "register"; "reinterpret_cast"; "return" - ; "short"; "signed"; "sizeof"; "static"; "static_assert"; "static_cast" - ; "struct"; "switch"; "template"; "this"; "thread_local"; "throw"; "true" - ; "try"; "typedef"; "typeid"; "typename"; "union"; "unsigned"; "using" - ; "virtual"; "void"; "volatile"; "wchar_t"; "while"; "xor"; "xor_eq" ] + [ "true" + ; "false" + ; "repeat" + ; "until" + ; "then" + ; "var" + ; "fvar" + ; "STAN_MAJOR" + ; "STAN_MINOR" + ; "STAN_PATCH" + ; "STAN_MATH_MAJOR" + ; "STAN_MATH_MINOR" + ; "STAN_MATH_PATCH" + ; "alignas" + ; "alignof" + ; "and" + ; "and_eq" + ; "asm" + ; "auto" + ; "bitand" + ; "bitor" + ; "bool" + ; "break" + ; "case" + ; "catch" + ; "char" + ; "char16_t" + ; "char32_t" + ; "class" + ; "compl" + ; "const" + ; "constexpr" + ; "const_cast" + ; "continue" + ; "decltype" + ; "default" + ; "delete" + ; "do" + ; "double" + ; "dynamic_cast" + ; "else" + ; "enum" + ; "explicit" + ; "export" + ; "extern" + ; "false" + ; "float" + ; "for" + ; "friend" + ; "goto" + ; "if" + ; "inline" + ; "int" + ; "long" + ; "mutable" + ; "namespace" + ; "new" + ; "noexcept" + ; "not" + ; "not_eq" + ; "nullptr" + ; "operator" + ; "or" + ; "or_eq" + ; "private" + ; "protected" + ; "public" + ; "register" + ; "reinterpret_cast" + ; "return" + ; "short" + ; "signed" + ; "sizeof" + ; "static" + ; "static_assert" + ; "static_cast" + ; "struct" + ; "switch" + ; "template" + ; "this" + ; "thread_local" + ; "throw" + ; "true" + ; "try" + ; "typedef" + ; "typeid" + ; "typename" + ; "union" + ; "unsigned" + ; "using" + ; "virtual" + ; "void" + ; "volatile" + ; "wchar_t" + ; "while" + ; "xor" + ; "xor_eq" + ] +;; let semantic_check_identifier id = Validate.( - if id.name = !model_name then - Semantic_error.ident_is_model_name id.id_loc id.name |> error - else if - String.is_suffix id.name ~suffix:"__" - || List.exists ~f:(fun str -> str = id.name) reserved_keywords + if id.name = !model_name + then Semantic_error.ident_is_model_name id.id_loc id.name |> error + else if String.is_suffix id.name ~suffix:"__" + || List.exists ~f:(fun str -> str = id.name) reserved_keywords then Semantic_error.ident_is_keyword id.id_loc id.name |> error else ok ()) +;; (* -- Operators ------------------------------------------------------------- *) let semantic_check_operator _ = Validate.ok () (* == Expressions =========================================================== *) -let arg_type x = (x.emeta.ad_level, x.emeta.type_) +let arg_type x = x.emeta.ad_level, x.emeta.type_ let get_arg_types = List.map ~f:arg_type (* -- Function application -------------------------------------------------- *) let semantic_check_fn_map_rect ~loc id es = Validate.( - match (id.name, es) with - | "map_rect", {expr= Variable arg1; _} :: _ + match id.name, es with + | "map_rect", { expr = Variable arg1; _ } :: _ when String.( - is_suffix arg1.name ~suffix:"_lp" - || is_suffix arg1.name ~suffix:"_rng") -> - Semantic_error.invalid_map_rect_fn loc arg1.name |> error + is_suffix arg1.name ~suffix:"_lp" || is_suffix arg1.name ~suffix:"_rng") -> + Semantic_error.invalid_map_rect_fn loc arg1.name |> error | _ -> ok ()) +;; let semantic_check_fn_conditioning ~loc id = Validate.( - if - List.exists ["_lpdf"; "_lpmf"; "_lcdf"; "_lccdf"] ~f:(fun x -> - String.is_suffix id.name ~suffix:x ) + if List.exists [ "_lpdf"; "_lpmf"; "_lcdf"; "_lccdf" ] ~f:(fun x -> + String.is_suffix id.name ~suffix:x) then Semantic_error.conditioning_required loc |> error else ok ()) +;; (** `Target+=` can only be used in model and functions with right suffix (same for tilde etc) *) let semantic_check_fn_target_plus_equals cf ~loc id = Validate.( - if - String.is_suffix id.name ~suffix:"_lp" - && not (cf.in_lp_fun_def || cf.current_block = Model) + if String.is_suffix id.name ~suffix:"_lp" + && not (cf.in_lp_fun_def || cf.current_block = Model) then Semantic_error.target_plusequals_outisde_model_or_logprob loc |> error else ok ()) +;; (** Rng functions cannot be used in Tp or Model and only in function defs with the right suffix *) let semantic_check_fn_rng cf ~loc id = Validate.( - if String.is_suffix id.name ~suffix:"_rng" && cf.in_toplevel_decl then - Semantic_error.invalid_decl_rng_fn loc |> error - else if - String.is_suffix id.name ~suffix:"_rng" - && ( (cf.in_fun_def && not cf.in_rng_fun_def) - || cf.current_block = TParam || cf.current_block = Model ) + if String.is_suffix id.name ~suffix:"_rng" && cf.in_toplevel_decl + then Semantic_error.invalid_decl_rng_fn loc |> error + else if String.is_suffix id.name ~suffix:"_rng" + && ((cf.in_fun_def && not cf.in_rng_fun_def) + || cf.current_block = TParam + || cf.current_block = Model) then Semantic_error.invalid_rng_fn loc |> error else ok ()) +;; let mk_fun_app ~is_cond_dist (x, y, z) = if is_cond_dist then CondDistApp (x, y, z) else FunApp (x, y, z) +;; (* Regular function application *) let semantic_check_fn_normal ~is_cond_dist ~loc id es = Validate.( match Symbol_table.look vm id.name with | Some (_, UnsizedType.UFun (_, Void)) -> - Semantic_error.returning_fn_expected_nonreturning_found loc id.name - |> error + Semantic_error.returning_fn_expected_nonreturning_found loc id.name |> error | Some (_, UFun (listedtypes, rt)) when not - (UnsizedType.check_compatible_arguments_mod_conv id.name - listedtypes (get_arg_types es)) -> - es - |> List.map ~f:type_of_expr_typed - |> Semantic_error.illtyped_userdefined_fn_app loc id.name listedtypes - rt - |> error + (UnsizedType.check_compatible_arguments_mod_conv + id.name + listedtypes + (get_arg_types es)) -> + es + |> List.map ~f:type_of_expr_typed + |> Semantic_error.illtyped_userdefined_fn_app loc id.name listedtypes rt + |> error | Some (_, UFun (_, ReturnType ut)) -> - mk_typed_expression - ~expr:(mk_fun_app ~is_cond_dist (UserDefined, id, es)) - ~ad_level:(expr_ad_lub es) ~type_:ut ~loc - |> ok + mk_typed_expression + ~expr:(mk_fun_app ~is_cond_dist (UserDefined, id, es)) + ~ad_level:(expr_ad_lub es) + ~type_:ut + ~loc + |> ok | Some _ -> - (* Check that Funaps are actually functions *) - Semantic_error.returning_fn_expected_nonfn_found loc id.name |> error + (* Check that Funaps are actually functions *) + Semantic_error.returning_fn_expected_nonfn_found loc id.name |> error | None -> - Semantic_error.returning_fn_expected_undeclaredident_found loc id.name - |> error) + Semantic_error.returning_fn_expected_undeclaredident_found loc id.name |> error) +;; (* Stan-Math function application *) let semantic_check_fn_stan_math ~is_cond_dist ~loc id es = - match - Stan_math_signatures.stan_math_returntype id.name (get_arg_types es) - with + match Stan_math_signatures.stan_math_returntype id.name (get_arg_types es) with | Some UnsizedType.Void -> - Semantic_error.returning_fn_expected_nonreturning_found loc id.name - |> Validate.error + Semantic_error.returning_fn_expected_nonreturning_found loc id.name |> Validate.error | Some (UnsizedType.ReturnType ut) -> - mk_typed_expression - ~expr:(mk_fun_app ~is_cond_dist (StanLib, id, es)) - ~ad_level:(expr_ad_lub es) ~type_:ut ~loc - |> Validate.ok + mk_typed_expression + ~expr:(mk_fun_app ~is_cond_dist (StanLib, id, es)) + ~ad_level:(expr_ad_lub es) + ~type_:ut + ~loc + |> Validate.ok | _ -> - es - |> List.map ~f:(fun e -> e.emeta.type_) - |> Semantic_error.illtyped_stanlib_fn_app loc id.name - |> Validate.error + es + |> List.map ~f:(fun e -> e.emeta.type_) + |> Semantic_error.illtyped_stanlib_fn_app loc id.name + |> Validate.error +;; let semantic_check_reduce_sum ~is_cond_dist ~loc id es = let arg_match (x_ad, x_t) y = @@ -303,46 +400,61 @@ let semantic_check_reduce_sum ~is_cond_dist ~loc id es = List.length a = List.length b && List.for_all2_exn ~f:arg_match a b in match es with - | { emeta= - { type_= + | { emeta = + { type_ = UnsizedType.UFun ( ((_, sliced_arg_fun_type) as sliced_arg_fun) :: (_, UInt) :: (_, UInt) :: fun_args - , ReturnType UReal ); _ }; _ } - :: sliced :: {emeta= {type_= UInt; _}; _} :: args + , ReturnType UReal ) + ; _ + } + ; _ + } + :: sliced :: { emeta = { type_ = UInt; _ }; _ } :: args when arg_match sliced_arg_fun sliced - && List.mem Stan_math_signatures.reduce_sum_slice_types - sliced.emeta.type_ ~equal:( = ) - && List.mem Stan_math_signatures.reduce_sum_slice_types - sliced_arg_fun_type ~equal:( = ) -> - if args_match fun_args args then - mk_typed_expression - ~expr:(mk_fun_app ~is_cond_dist (StanLib, id, es)) - ~ad_level:(expr_ad_lub es) ~type_:UnsizedType.UReal ~loc - |> Validate.ok - else - Semantic_error.illtyped_reduce_sum loc id.name - (List.map ~f:type_of_expr_typed es) - (sliced_arg_fun :: fun_args) - |> Validate.error - | _ -> - es - |> List.map ~f:type_of_expr_typed - |> Semantic_error.illtyped_reduce_sum_generic loc id.name + && List.mem + Stan_math_signatures.reduce_sum_slice_types + sliced.emeta.type_ + ~equal:( = ) + && List.mem + Stan_math_signatures.reduce_sum_slice_types + sliced_arg_fun_type + ~equal:( = ) -> + if args_match fun_args args + then + mk_typed_expression + ~expr:(mk_fun_app ~is_cond_dist (StanLib, id, es)) + ~ad_level:(expr_ad_lub es) + ~type_:UnsizedType.UReal + ~loc + |> Validate.ok + else + Semantic_error.illtyped_reduce_sum + loc + id.name + (List.map ~f:type_of_expr_typed es) + (sliced_arg_fun :: fun_args) |> Validate.error + | _ -> + es + |> List.map ~f:type_of_expr_typed + |> Semantic_error.illtyped_reduce_sum_generic loc id.name + |> Validate.error +;; let fn_kind_from_application id es = (* We need to check an application here, rather than a mere name of the function because, technically, user defined functions can shadow constants in StanLib. *) - if - Stan_math_signatures.stan_math_returntype id.name - (List.map ~f:(fun x -> (x.emeta.ad_level, x.emeta.type_)) es) - <> None - || Symbol_table.look vm id.name = None - && Stan_math_signatures.is_stan_math_function_name id.name + if Stan_math_signatures.stan_math_returntype + id.name + (List.map ~f:(fun x -> x.emeta.ad_level, x.emeta.type_) es) + <> None + || (Symbol_table.look vm id.name = None + && Stan_math_signatures.is_stan_math_function_name id.name) then StanLib else UserDefined +;; (** Determines the function kind based on the identifier and performs the corresponding semantic check @@ -350,116 +462,130 @@ let fn_kind_from_application id es = let semantic_check_fn ~is_cond_dist ~loc id es = match fn_kind_from_application id es with | StanLib when Stan_math_signatures.is_reduce_sum_fn id.name -> - semantic_check_reduce_sum ~is_cond_dist ~loc id es + semantic_check_reduce_sum ~is_cond_dist ~loc id es | StanLib -> semantic_check_fn_stan_math ~is_cond_dist ~loc id es | UserDefined -> semantic_check_fn_normal ~is_cond_dist ~loc id es +;; (* -- Ternary If ------------------------------------------------------------ *) let semantic_check_ternary_if loc (pe, te, fe) = Validate.( let err = - Semantic_error.illtyped_ternary_if loc pe.emeta.type_ te.emeta.type_ - fe.emeta.type_ + Semantic_error.illtyped_ternary_if loc pe.emeta.type_ te.emeta.type_ fe.emeta.type_ in - if pe.emeta.type_ = UInt then + if pe.emeta.type_ = UInt + then ( match UnsizedType.common_type (te.emeta.type_, fe.emeta.type_) with | Some type_ -> - mk_typed_expression - ~expr:(TernaryIf (pe, te, fe)) - ~ad_level:(expr_ad_lub [pe; te; fe]) - ~type_ ~loc - |> ok - | None -> error err + mk_typed_expression + ~expr:(TernaryIf (pe, te, fe)) + ~ad_level:(expr_ad_lub [ pe; te; fe ]) + ~type_ + ~loc + |> ok + | None -> error err) else error err) +;; (* -- Binary (Infix) Operators ---------------------------------------------- *) let semantic_check_binop loc op (le, re) = Validate.( - let err = - Semantic_error.illtyped_binary_op loc op le.emeta.type_ re.emeta.type_ - in - [le; re] |> List.map ~f:arg_type + let err = Semantic_error.illtyped_binary_op loc op le.emeta.type_ re.emeta.type_ in + [ le; re ] + |> List.map ~f:arg_type |> Stan_math_signatures.operator_stan_math_return_type op |> Option.value_map ~default:(error err) ~f:(function - | ReturnType type_ -> + | ReturnType type_ -> mk_typed_expression ~expr:(BinOp (le, op, re)) - ~ad_level:(expr_ad_lub [le; re]) - ~type_ ~loc + ~ad_level:(expr_ad_lub [ le; re ]) + ~type_ + ~loc |> ok - | Void -> error err )) + | Void -> error err)) +;; let to_exn v = - v |> Validate.to_result + v + |> Validate.to_result |> Result.map_error ~f:Fmt.(to_to_string @@ list ~sep:cut Semantic_error.pp) |> Result.ok_or_failwith +;; let semantic_check_binop_exn loc op (le, re) = semantic_check_binop loc op (le, re) |> to_exn +;; (* -- Prefix Operators ------------------------------------------------------ *) let semantic_check_prefixop loc op e = Validate.( let err = Semantic_error.illtyped_prefix_op loc op e.emeta.type_ in - Stan_math_signatures.operator_stan_math_return_type op [arg_type e] + Stan_math_signatures.operator_stan_math_return_type op [ arg_type e ] |> Option.value_map ~default:(error err) ~f:(function - | ReturnType type_ -> + | ReturnType type_ -> mk_typed_expression ~expr:(PrefixOp (op, e)) - ~ad_level:(expr_ad_lub [e]) - ~type_ ~loc + ~ad_level:(expr_ad_lub [ e ]) + ~type_ + ~loc |> ok - | Void -> error err )) + | Void -> error err)) +;; (* -- Postfix operators ----------------------------------------------------- *) let semantic_check_postfixop loc op e = Validate.( let err = Semantic_error.illtyped_postfix_op loc op e.emeta.type_ in - Stan_math_signatures.operator_stan_math_return_type op [arg_type e] + Stan_math_signatures.operator_stan_math_return_type op [ arg_type e ] |> Option.value_map ~default:(error err) ~f:(function - | ReturnType type_ -> + | ReturnType type_ -> mk_typed_expression ~expr:(PostfixOp (e, op)) - ~ad_level:(expr_ad_lub [e]) - ~type_ ~loc + ~ad_level:(expr_ad_lub [ e ]) + ~type_ + ~loc |> ok - | Void -> error err )) + | Void -> error err)) +;; (* -- Variables ------------------------------------------------------------- *) let semantic_check_variable cf loc id = Validate.( match Symbol_table.look vm id.name with - | None when not (Stan_math_signatures.is_stan_math_function_name id.name) - -> - Semantic_error.ident_not_in_scope loc id.name |> error + | None when not (Stan_math_signatures.is_stan_math_function_name id.name) -> + Semantic_error.ident_not_in_scope loc id.name |> error | None -> - mk_typed_expression ~expr:(Variable id) - ~ad_level: - (calculate_autodifftype cf MathLibrary UMathLibraryFunction) - ~type_:UMathLibraryFunction ~loc - |> ok + mk_typed_expression + ~expr:(Variable id) + ~ad_level:(calculate_autodifftype cf MathLibrary UMathLibraryFunction) + ~type_:UMathLibraryFunction + ~loc + |> ok | Some ((Param | TParam | GQuant), _) when cf.in_toplevel_decl -> - Semantic_error.non_data_variable_size_decl loc |> error + Semantic_error.non_data_variable_size_decl loc |> error | Some (originblock, type_) -> - mk_typed_expression ~expr:(Variable id) - ~ad_level:(calculate_autodifftype cf originblock type_) - ~type_ ~loc - |> ok) + mk_typed_expression + ~expr:(Variable id) + ~ad_level:(calculate_autodifftype cf originblock type_) + ~type_ + ~loc + |> ok) +;; (* -- Conditioned Distribution Application ---------------------------------- *) let semantic_check_conddist_name ~loc id = Validate.( - if - List.exists - ~f:(fun x -> String.is_suffix id.name ~suffix:x) - ["_lpdf"; "_lpmf"; "_lcdf"; "_lccdf"] + if List.exists + ~f:(fun x -> String.is_suffix id.name ~suffix:x) + [ "_lpdf"; "_lpmf"; "_lcdf"; "_lccdf" ] then ok () else Semantic_error.conditional_notation_not_allowed loc |> error) +;; (* -- Array Expressions ----------------------------------------------------- *) @@ -467,108 +593,107 @@ let check_consistent_types ad_level type_ es = let f state e = match state with | Error e -> Error e - | Ok (ad, ty) -> ( - let ad = - if UnsizedType.autodifftype_can_convert e.emeta.ad_level ad then - e.emeta.ad_level - else ad - in - match UnsizedType.common_type (ty, e.emeta.type_) with - | Some ty -> Ok (ad, ty) - | None -> Error (ty, e.emeta) ) + | Ok (ad, ty) -> + let ad = + if UnsizedType.autodifftype_can_convert e.emeta.ad_level ad + then e.emeta.ad_level + else ad + in + (match UnsizedType.common_type (ty, e.emeta.type_) with + | Some ty -> Ok (ad, ty) + | None -> Error (ty, e.emeta)) in List.fold ~init:(Ok (ad_level, type_)) ~f es +;; let semantic_check_array_expr ~loc es = Validate.( match es with | [] -> Semantic_error.empty_array loc |> error - | {emeta= {ad_level; type_; _}; _} :: elements -> ( - match check_consistent_types ad_level type_ elements with + | { emeta = { ad_level; type_; _ }; _ } :: elements -> + (match check_consistent_types ad_level type_ elements with | Error (ty, meta) -> - Semantic_error.mismatched_array_types meta.loc ty meta.type_ |> error + Semantic_error.mismatched_array_types meta.loc ty meta.type_ |> error | Ok (ad_level, type_) -> - let type_ = UnsizedType.UArray type_ in - mk_typed_expression ~expr:(ArrayExpr es) ~ad_level ~type_ ~loc |> ok - )) + let type_ = UnsizedType.UArray type_ in + mk_typed_expression ~expr:(ArrayExpr es) ~ad_level ~type_ ~loc |> ok)) +;; (* -- Row Vector Expresssion ------------------------------------------------ *) let semantic_check_rowvector ~loc es = Validate.( match es with - | {emeta= {ad_level; type_= UnsizedType.URowVector; _}; _} :: elements -> ( - match check_consistent_types ad_level URowVector elements with + | { emeta = { ad_level; type_ = UnsizedType.URowVector; _ }; _ } :: elements -> + (match check_consistent_types ad_level URowVector elements with | Ok (ad_level, _) -> - mk_typed_expression ~expr:(RowVectorExpr es) ~ad_level ~type_:UMatrix - ~loc - |> ok + mk_typed_expression ~expr:(RowVectorExpr es) ~ad_level ~type_:UMatrix ~loc |> ok | Error (_, meta) -> - Semantic_error.invalid_matrix_types meta.loc meta.type_ |> error ) - | _ -> ( - match check_consistent_types DataOnly UReal es with + Semantic_error.invalid_matrix_types meta.loc meta.type_ |> error) + | _ -> + (match check_consistent_types DataOnly UReal es with | Ok (ad_level, _) -> - mk_typed_expression ~expr:(RowVectorExpr es) ~ad_level - ~type_:URowVector ~loc - |> ok + mk_typed_expression ~expr:(RowVectorExpr es) ~ad_level ~type_:URowVector ~loc + |> ok | Error (_, meta) -> - Semantic_error.invalid_row_vector_types meta.loc meta.type_ |> error - )) + Semantic_error.invalid_row_vector_types meta.loc meta.type_ |> error)) +;; (* -- Indexed Expressions --------------------------------------------------- *) -let tuple2 a b = (a, b) -let tuple3 a b c = (a, b, c) +let tuple2 a b = a, b +let tuple3 a b c = a, b, c let index_with_type idx = match idx with - | Single e -> (idx, e.emeta.type_) - | _ -> (idx, UnsizedType.UInt) + | Single e -> idx, e.emeta.type_ + | _ -> idx, UnsizedType.UInt +;; let inferred_unsizedtype_of_indexed ~loc ut indices = let rec aux k ut xs = - match (ut, xs) with - | UnsizedType.UMatrix, [(All, _); (Single _, UnsizedType.UInt)] - |UMatrix, [(Upfrom _, _); (Single _, UInt)] - |UMatrix, [(Downfrom _, _); (Single _, UInt)] - |UMatrix, [(Between _, _); (Single _, UInt)] - |UMatrix, [(Single _, UArray UInt); (Single _, UInt)] -> - k @@ Validate.ok UnsizedType.UVector + match ut, xs with + | UnsizedType.UMatrix, [ (All, _); (Single _, UnsizedType.UInt) ] + | UMatrix, [ (Upfrom _, _); (Single _, UInt) ] + | UMatrix, [ (Downfrom _, _); (Single _, UInt) ] + | UMatrix, [ (Between _, _); (Single _, UInt) ] + | UMatrix, [ (Single _, UArray UInt); (Single _, UInt) ] -> + k @@ Validate.ok UnsizedType.UVector | _, [] -> k @@ Validate.ok ut - | _, next :: rest -> ( - match next with - | Single _, UInt -> ( - match ut with + | _, next :: rest -> + (match next with + | Single _, UInt -> + (match ut with | UArray inner_ty -> aux k inner_ty rest | UVector | URowVector -> aux k UReal rest | UMatrix -> aux k URowVector rest - | _ -> Semantic_error.not_indexable loc ut |> Validate.error ) - | _ -> ( - match ut with + | _ -> Semantic_error.not_indexable loc ut |> Validate.error) + | _ -> + (match ut with | UArray inner_ty -> - let k' = - Fn.compose k (Validate.map ~f:(fun t -> UnsizedType.UArray t)) - in - aux k' inner_ty rest + let k' = Fn.compose k (Validate.map ~f:(fun t -> UnsizedType.UArray t)) in + aux k' inner_ty rest | UVector | URowVector | UMatrix -> aux k ut rest - | _ -> Semantic_error.not_indexable loc ut |> Validate.error ) ) + | _ -> Semantic_error.not_indexable loc ut |> Validate.error)) in aux Fn.id ut (List.map ~f:index_with_type indices) +;; let inferred_unsizedtype_of_indexed_exn ~loc ut indices = inferred_unsizedtype_of_indexed ~loc ut indices |> to_exn +;; let inferred_ad_type_of_indexed at uindices = UnsizedType.lub_ad_type - ( at + (at :: List.map ~f:(function | All -> UnsizedType.DataOnly | Single ue1 | Upfrom ue1 | Downfrom ue1 -> - UnsizedType.lub_ad_type [at; ue1.emeta.ad_level] + UnsizedType.lub_ad_type [ at; ue1.emeta.ad_level ] | Between (ue1, ue2) -> - UnsizedType.lub_ad_type - [at; ue1.emeta.ad_level; ue2.emeta.ad_level]) - uindices ) + UnsizedType.lub_ad_type [ at; ue1.emeta.ad_level; ue2.emeta.ad_level ]) + uindices) +;; let rec semantic_check_indexed ~loc ~cf e indices = Validate.( @@ -581,154 +706,165 @@ let rec semantic_check_indexed ~loc ~cf e indices = uindices |> inferred_unsizedtype_of_indexed ~loc ue.emeta.type_ |> map ~f:(fun ut -> - mk_typed_expression - ~expr:(Indexed (ue, uindices)) - ~ad_level:at ~type_:ut ~loc )) + mk_typed_expression ~expr:(Indexed (ue, uindices)) ~ad_level:at ~type_:ut ~loc)) and semantic_check_index cf = function | All -> Validate.ok All (* Check that indexes have int (container) type *) | Single e -> - Validate.( - semantic_check_expression cf e - >>= fun ue -> - if has_int_type ue || has_int_array_type ue then ok @@ Single ue - else - Semantic_error.int_intarray_or_range_expected ue.emeta.loc - ue.emeta.type_ - |> error) + Validate.( + semantic_check_expression cf e + >>= fun ue -> + if has_int_type ue || has_int_array_type ue + then ok @@ Single ue + else + Semantic_error.int_intarray_or_range_expected ue.emeta.loc ue.emeta.type_ |> error) | Upfrom e -> - semantic_check_expression_of_int_type cf e "Range bound" - |> Validate.map ~f:(fun e -> Upfrom e) + semantic_check_expression_of_int_type cf e "Range bound" + |> Validate.map ~f:(fun e -> Upfrom e) | Downfrom e -> - semantic_check_expression_of_int_type cf e "Range bound" - |> Validate.map ~f:(fun e -> Downfrom e) + semantic_check_expression_of_int_type cf e "Range bound" + |> Validate.map ~f:(fun e -> Downfrom e) | Between (e1, e2) -> - let le = semantic_check_expression_of_int_type cf e1 "Range bound" - and ue = semantic_check_expression_of_int_type cf e2 "Range bound" in - Validate.liftA2 (fun l u -> Between (l, u)) le ue + let le = semantic_check_expression_of_int_type cf e1 "Range bound" + and ue = semantic_check_expression_of_int_type cf e2 "Range bound" in + Validate.liftA2 (fun l u -> Between (l, u)) le ue (* -- Top-level expressions ------------------------------------------------- *) -and semantic_check_expression cf ({emeta; expr} : Ast.untyped_expression) : - Ast.typed_expression Validate.t = +and semantic_check_expression cf ({ emeta; expr } : Ast.untyped_expression) + : Ast.typed_expression Validate.t + = match expr with | TernaryIf (e1, e2, e3) -> - let pe = semantic_check_expression cf e1 - and te = semantic_check_expression cf e2 - and fe = semantic_check_expression cf e3 in - Validate.(liftA3 tuple3 pe te fe >>= semantic_check_ternary_if emeta.loc) + let pe = semantic_check_expression cf e1 + and te = semantic_check_expression cf e2 + and fe = semantic_check_expression cf e3 in + Validate.(liftA3 tuple3 pe te fe >>= semantic_check_ternary_if emeta.loc) | BinOp (e1, op, e2) -> - let le = semantic_check_expression cf e1 - and re = semantic_check_expression cf e2 - and warn_int_division (x, y) = - match (x.emeta.type_, y.emeta.type_, op) with - | UInt, UInt, Divide -> - let hint ppf () = - match (x.expr, y.expr) with - | IntNumeral x, _ -> - Fmt.pf ppf "%s.0 / %a" x Pretty_printing.pp_expression y - | _, Ast.IntNumeral y -> - Fmt.pf ppf "%a / %s.0" Pretty_printing.pp_expression x y - | _ -> - Fmt.pf ppf "%a * 1.0 / %a" Pretty_printing.pp_expression x - Pretty_printing.pp_expression y - in - Fmt.pr - "@[@[Info: Found int division at %s:@]@ @[%a@]@,@[%a@]@ @[%a@]@,@[%a@]@]" - (Location_span.to_string x.emeta.loc) - Pretty_printing.pp_expression {expr; emeta} Fmt.text - "Values will be rounded towards zero. If rounding is not \ - desired you can write the division as" - hint () Fmt.text - "If rounding is intended please use the integer division \ - operator %/%." ; - (x, y) - | _ -> (x, y) - in - Validate.( - liftA2 tuple2 le re |> map ~f:warn_int_division - |> apply_const (semantic_check_operator op) - >>= semantic_check_binop emeta.loc op) + let le = semantic_check_expression cf e1 + and re = semantic_check_expression cf e2 + and warn_int_division (x, y) = + match x.emeta.type_, y.emeta.type_, op with + | UInt, UInt, Divide -> + let hint ppf () = + match x.expr, y.expr with + | IntNumeral x, _ -> Fmt.pf ppf "%s.0 / %a" x Pretty_printing.pp_expression y + | _, Ast.IntNumeral y -> + Fmt.pf ppf "%a / %s.0" Pretty_printing.pp_expression x y + | _ -> + Fmt.pf + ppf + "%a * 1.0 / %a" + Pretty_printing.pp_expression + x + Pretty_printing.pp_expression + y + in + Fmt.pr + "@[@[Info: Found int division at %s:@]@ @[%a@]@,\ + @[%a@]@ @[%a@]@,\ + @[%a@]@]" + (Location_span.to_string x.emeta.loc) + Pretty_printing.pp_expression + { expr; emeta } + Fmt.text + "Values will be rounded towards zero. If rounding is not desired you can write \ + the division as" + hint + () + Fmt.text + "If rounding is intended please use the integer division operator %/%."; + x, y + | _ -> x, y + in + Validate.( + liftA2 tuple2 le re + |> map ~f:warn_int_division + |> apply_const (semantic_check_operator op) + >>= semantic_check_binop emeta.loc op) | PrefixOp (op, e) -> - Validate.( - semantic_check_expression cf e - |> apply_const (semantic_check_operator op) - >>= semantic_check_prefixop emeta.loc op) + Validate.( + semantic_check_expression cf e + |> apply_const (semantic_check_operator op) + >>= semantic_check_prefixop emeta.loc op) | PostfixOp (e, op) -> - Validate.( - semantic_check_expression cf e - |> apply_const (semantic_check_operator op) - >>= semantic_check_postfixop emeta.loc op) + Validate.( + semantic_check_expression cf e + |> apply_const (semantic_check_operator op) + >>= semantic_check_postfixop emeta.loc op) | Variable id -> - semantic_check_variable cf emeta.loc id - |> Validate.apply_const (semantic_check_identifier id) - | IntNumeral s -> ( - match float_of_string_opt s with + semantic_check_variable cf emeta.loc id + |> Validate.apply_const (semantic_check_identifier id) + | IntNumeral s -> + (match float_of_string_opt s with | Some i when i < 2_147_483_648.0 -> - mk_typed_expression ~expr:(IntNumeral s) ~ad_level:DataOnly ~type_:UInt - ~loc:emeta.loc - |> Validate.ok - | _ -> Semantic_error.bad_int_literal emeta.loc |> Validate.error ) - | RealNumeral s -> - mk_typed_expression ~expr:(RealNumeral s) ~ad_level:DataOnly ~type_:UReal + mk_typed_expression + ~expr:(IntNumeral s) + ~ad_level:DataOnly + ~type_:UInt ~loc:emeta.loc |> Validate.ok - | FunApp (_, id, es) -> - semantic_check_funapp ~is_cond_dist:false id es cf emeta - | CondDistApp (_, id, es) -> - semantic_check_funapp ~is_cond_dist:true id es cf emeta + | _ -> Semantic_error.bad_int_literal emeta.loc |> Validate.error) + | RealNumeral s -> + mk_typed_expression + ~expr:(RealNumeral s) + ~ad_level:DataOnly + ~type_:UReal + ~loc:emeta.loc + |> Validate.ok + | FunApp (_, id, es) -> semantic_check_funapp ~is_cond_dist:false id es cf emeta + | CondDistApp (_, id, es) -> semantic_check_funapp ~is_cond_dist:true id es cf emeta | GetLP -> - (* Target+= can only be used in model and functions with right suffix (same for tilde etc) *) - if - not - ( cf.in_lp_fun_def || cf.current_block = Model - || cf.current_block = TParam ) - then - Semantic_error.target_plusequals_outisde_model_or_logprob emeta.loc - |> Validate.error - else - mk_typed_expression ~expr:GetLP - ~ad_level:(calculate_autodifftype cf cf.current_block UReal) - ~type_:UReal ~loc:emeta.loc - |> Validate.ok + (* Target+= can only be used in model and functions with right suffix (same for tilde etc) *) + if not (cf.in_lp_fun_def || cf.current_block = Model || cf.current_block = TParam) + then + Semantic_error.target_plusequals_outisde_model_or_logprob emeta.loc + |> Validate.error + else + mk_typed_expression + ~expr:GetLP + ~ad_level:(calculate_autodifftype cf cf.current_block UReal) + ~type_:UReal + ~loc:emeta.loc + |> Validate.ok | GetTarget -> - (* Target+= can only be used in model and functions with right suffix (same for tilde etc) *) - if - not - ( cf.in_lp_fun_def || cf.current_block = Model - || cf.current_block = TParam ) - then - Semantic_error.target_plusequals_outisde_model_or_logprob emeta.loc - |> Validate.error - else - mk_typed_expression ~expr:GetTarget - ~ad_level:(calculate_autodifftype cf cf.current_block UReal) - ~type_:UReal ~loc:emeta.loc - |> Validate.ok + (* Target+= can only be used in model and functions with right suffix (same for tilde etc) *) + if not (cf.in_lp_fun_def || cf.current_block = Model || cf.current_block = TParam) + then + Semantic_error.target_plusequals_outisde_model_or_logprob emeta.loc + |> Validate.error + else + mk_typed_expression + ~expr:GetTarget + ~ad_level:(calculate_autodifftype cf cf.current_block UReal) + ~type_:UReal + ~loc:emeta.loc + |> Validate.ok | ArrayExpr es -> - Validate.( - es - |> List.map ~f:(semantic_check_expression cf) - |> sequence - >>= fun ues -> semantic_check_array_expr ~loc:emeta.loc ues) + Validate.( + es + |> List.map ~f:(semantic_check_expression cf) + |> sequence + >>= fun ues -> semantic_check_array_expr ~loc:emeta.loc ues) | RowVectorExpr es -> - Validate.( - es - |> List.map ~f:(semantic_check_expression cf) - |> sequence - >>= semantic_check_rowvector ~loc:emeta.loc) + Validate.( + es + |> List.map ~f:(semantic_check_expression cf) + |> sequence + >>= semantic_check_rowvector ~loc:emeta.loc) | Paren e -> - semantic_check_expression cf e - |> Validate.map ~f:(fun ue -> - mk_typed_expression ~expr:(Paren ue) ~ad_level:ue.emeta.ad_level - ~type_:ue.emeta.type_ ~loc:emeta.loc ) + semantic_check_expression cf e + |> Validate.map ~f:(fun ue -> + mk_typed_expression + ~expr:(Paren ue) + ~ad_level:ue.emeta.ad_level + ~type_:ue.emeta.type_ + ~loc:emeta.loc) | Indexed (e, indices) -> semantic_check_indexed ~loc:emeta.loc ~cf e indices and semantic_check_funapp ~is_cond_dist id es cf emeta = let name_check = - if is_cond_dist then semantic_check_conddist_name - else semantic_check_fn_conditioning + if is_cond_dist then semantic_check_conddist_name else semantic_check_fn_conditioning in Validate.( es @@ -746,77 +882,72 @@ and semantic_check_expression_of_int_type cf e name = Validate.( semantic_check_expression cf e >>= fun ue -> - if has_int_type ue then ok ue + if has_int_type ue + then ok ue else Semantic_error.int_expected ue.emeta.loc name ue.emeta.type_ |> error) and semantic_check_expression_of_int_or_real_type cf e name = Validate.( semantic_check_expression cf e >>= fun ue -> - if has_int_or_real_type ue then ok ue - else - Semantic_error.int_or_real_expected ue.emeta.loc name ue.emeta.type_ - |> error) + if has_int_or_real_type ue + then ok ue + else Semantic_error.int_or_real_expected ue.emeta.loc name ue.emeta.type_ |> error) +;; let semantic_check_expression_of_scalar_or_type cf t e name = Validate.( semantic_check_expression cf e >>= fun ue -> - if UnsizedType.is_scalar_type ue.emeta.type_ || ue.emeta.type_ = t then - ok ue + if UnsizedType.is_scalar_type ue.emeta.type_ || ue.emeta.type_ = t + then ok ue else - Semantic_error.scalar_or_type_expected ue.emeta.loc name t ue.emeta.type_ - |> error) + Semantic_error.scalar_or_type_expected ue.emeta.loc name t ue.emeta.type_ |> error) +;; (* -- Sized Types ----------------------------------------------------------- *) let rec semantic_check_sizedtype cf = function | SizedType.SInt -> Validate.ok SizedType.SInt | SReal -> Validate.ok SizedType.SReal | SVector e -> - semantic_check_expression_of_int_type cf e "Vector sizes" - |> Validate.map ~f:(fun ue -> SizedType.SVector ue) + semantic_check_expression_of_int_type cf e "Vector sizes" + |> Validate.map ~f:(fun ue -> SizedType.SVector ue) | SRowVector e -> - semantic_check_expression_of_int_type cf e "Row vector sizes" - |> Validate.map ~f:(fun ue -> SizedType.SRowVector ue) + semantic_check_expression_of_int_type cf e "Row vector sizes" + |> Validate.map ~f:(fun ue -> SizedType.SRowVector ue) | SMatrix (e1, e2) -> - let ue1 = semantic_check_expression_of_int_type cf e1 "Matrix sizes" - and ue2 = semantic_check_expression_of_int_type cf e2 "Matrix sizes" in - Validate.liftA2 (fun ue1 ue2 -> SizedType.SMatrix (ue1, ue2)) ue1 ue2 + let ue1 = semantic_check_expression_of_int_type cf e1 "Matrix sizes" + and ue2 = semantic_check_expression_of_int_type cf e2 "Matrix sizes" in + Validate.liftA2 (fun ue1 ue2 -> SizedType.SMatrix (ue1, ue2)) ue1 ue2 | SArray (st, e) -> - let ust = semantic_check_sizedtype cf st - and ue = semantic_check_expression_of_int_type cf e "Array sizes" in - Validate.liftA2 (fun ust ue -> SizedType.SArray (ust, ue)) ust ue + let ust = semantic_check_sizedtype cf st + and ue = semantic_check_expression_of_int_type cf e "Array sizes" in + Validate.liftA2 (fun ust ue -> SizedType.SArray (ust, ue)) ust ue +;; (* -- Transformations ------------------------------------------------------- *) let semantic_check_transformation cf ut = function | Program.Identity -> Validate.ok Program.Identity | Lower e -> - semantic_check_expression_of_scalar_or_type cf ut e "Lower bound" - |> Validate.map ~f:(fun ue -> Program.Lower ue) + semantic_check_expression_of_scalar_or_type cf ut e "Lower bound" + |> Validate.map ~f:(fun ue -> Program.Lower ue) | Upper e -> - semantic_check_expression_of_scalar_or_type cf ut e "Upper bound" - |> Validate.map ~f:(fun ue -> Program.Upper ue) + semantic_check_expression_of_scalar_or_type cf ut e "Upper bound" + |> Validate.map ~f:(fun ue -> Program.Upper ue) | LowerUpper (e1, e2) -> - let ue1 = - semantic_check_expression_of_scalar_or_type cf ut e1 "Lower bound" - and ue2 = - semantic_check_expression_of_scalar_or_type cf ut e2 "Upper bound" - in - Validate.liftA2 (fun ue1 ue2 -> Program.LowerUpper (ue1, ue2)) ue1 ue2 + let ue1 = semantic_check_expression_of_scalar_or_type cf ut e1 "Lower bound" + and ue2 = semantic_check_expression_of_scalar_or_type cf ut e2 "Upper bound" in + Validate.liftA2 (fun ue1 ue2 -> Program.LowerUpper (ue1, ue2)) ue1 ue2 | Offset e -> - semantic_check_expression_of_scalar_or_type cf ut e "Offset" - |> Validate.map ~f:(fun ue -> Program.Offset ue) + semantic_check_expression_of_scalar_or_type cf ut e "Offset" + |> Validate.map ~f:(fun ue -> Program.Offset ue) | Multiplier e -> - semantic_check_expression_of_scalar_or_type cf ut e "Multiplier" - |> Validate.map ~f:(fun ue -> Program.Multiplier ue) + semantic_check_expression_of_scalar_or_type cf ut e "Multiplier" + |> Validate.map ~f:(fun ue -> Program.Multiplier ue) | OffsetMultiplier (e1, e2) -> - let ue1 = semantic_check_expression_of_scalar_or_type cf ut e1 "Offset" - and ue2 = - semantic_check_expression_of_scalar_or_type cf ut e2 "Multiplier" - in - Validate.liftA2 - (fun ue1 ue2 -> Program.OffsetMultiplier (ue1, ue2)) - ue1 ue2 + let ue1 = semantic_check_expression_of_scalar_or_type cf ut e1 "Offset" + and ue2 = semantic_check_expression_of_scalar_or_type cf ut e2 "Multiplier" in + Validate.liftA2 (fun ue1 ue2 -> Program.OffsetMultiplier (ue1, ue2)) ue1 ue2 | Ordered -> Validate.ok Program.Ordered | PositiveOrdered -> Validate.ok Program.PositiveOrdered | Simplex -> Validate.ok Program.Simplex @@ -825,38 +956,38 @@ let semantic_check_transformation cf ut = function | CholeskyCov -> Validate.ok Program.CholeskyCov | Correlation -> Validate.ok Program.Correlation | Covariance -> Validate.ok Program.Covariance +;; (* -- Printables ------------------------------------------------------------ *) let semantic_check_printable cf = function | PString s -> Validate.ok @@ PString s (* Print/reject expressions cannot be of function type. *) - | PExpr e -> ( - Validate.( - semantic_check_expression cf e - >>= fun ue -> - match ue.emeta.type_ with - | UFun _ | UMathLibraryFunction -> - Semantic_error.not_printable ue.emeta.loc |> error - | _ -> ok @@ PExpr ue) ) + | PExpr e -> + Validate.( + semantic_check_expression cf e + >>= fun ue -> + (match ue.emeta.type_ with + | UFun _ | UMathLibraryFunction -> + Semantic_error.not_printable ue.emeta.loc |> error + | _ -> ok @@ PExpr ue)) +;; (* -- Truncations ----------------------------------------------------------- *) let semantic_check_truncation cf = function | NoTruncate -> Validate.ok NoTruncate | TruncateUpFrom e -> - semantic_check_expression_of_int_or_real_type cf e "Truncation bound" - |> Validate.map ~f:(fun ue -> TruncateUpFrom ue) + semantic_check_expression_of_int_or_real_type cf e "Truncation bound" + |> Validate.map ~f:(fun ue -> TruncateUpFrom ue) | TruncateDownFrom e -> - semantic_check_expression_of_int_or_real_type cf e "Truncation bound" - |> Validate.map ~f:(fun ue -> TruncateDownFrom ue) + semantic_check_expression_of_int_or_real_type cf e "Truncation bound" + |> Validate.map ~f:(fun ue -> TruncateDownFrom ue) | TruncateBetween (e1, e2) -> - let ue1 = - semantic_check_expression_of_int_or_real_type cf e1 "Truncation bound" - and ue2 = - semantic_check_expression_of_int_or_real_type cf e2 "Truncation bound" - in - Validate.liftA2 (fun ue1 ue2 -> TruncateBetween (ue1, ue2)) ue1 ue2 + let ue1 = semantic_check_expression_of_int_or_real_type cf e1 "Truncation bound" + and ue2 = semantic_check_expression_of_int_or_real_type cf e2 "Truncation bound" in + Validate.liftA2 (fun ue1 ue2 -> TruncateBetween (ue1, ue2)) ue1 ue2 +;; (* == Statements ============================================================ *) @@ -864,62 +995,57 @@ let semantic_check_truncation cf = function let semantic_check_nrfn_target ~loc ~cf id = Validate.( - if - String.is_suffix id.name ~suffix:"_lp" - && not (cf.in_lp_fun_def || cf.current_block = Model) + if String.is_suffix id.name ~suffix:"_lp" + && not (cf.in_lp_fun_def || cf.current_block = Model) then Semantic_error.target_plusequals_outisde_model_or_logprob loc |> error else ok ()) +;; let semantic_check_nrfn_normal ~loc id es = Validate.( match Symbol_table.look vm id.name with | Some (_, UFun (listedtypes, Void)) - when UnsizedType.check_compatible_arguments_mod_conv id.name listedtypes + when UnsizedType.check_compatible_arguments_mod_conv + id.name + listedtypes (get_arg_types es) -> - mk_typed_statement - ~stmt:(NRFunApp (UserDefined, id, es)) - ~return_type:NoReturnType ~loc - |> ok + mk_typed_statement + ~stmt:(NRFunApp (UserDefined, id, es)) + ~return_type:NoReturnType + ~loc + |> ok | Some (_, UFun (listedtypes, Void)) -> - es - |> List.map ~f:type_of_expr_typed - |> Semantic_error.illtyped_userdefined_fn_app loc id.name listedtypes - Void - |> error + es + |> List.map ~f:type_of_expr_typed + |> Semantic_error.illtyped_userdefined_fn_app loc id.name listedtypes Void + |> error | Some (_, UFun (_, ReturnType _)) -> - Semantic_error.nonreturning_fn_expected_returning_found loc id.name - |> error - | Some _ -> - Semantic_error.nonreturning_fn_expected_nonfn_found loc id.name - |> error + Semantic_error.nonreturning_fn_expected_returning_found loc id.name |> error + | Some _ -> Semantic_error.nonreturning_fn_expected_nonfn_found loc id.name |> error | None -> - Semantic_error.nonreturning_fn_expected_undeclaredident_found loc - id.name - |> error) + Semantic_error.nonreturning_fn_expected_undeclaredident_found loc id.name |> error) +;; let semantic_check_nrfn_stan_math ~loc id es = Validate.( - match - Stan_math_signatures.stan_math_returntype id.name (get_arg_types es) - with + match Stan_math_signatures.stan_math_returntype id.name (get_arg_types es) with | Some UnsizedType.Void -> - mk_typed_statement - ~stmt:(NRFunApp (StanLib, id, es)) - ~return_type:NoReturnType ~loc - |> ok + mk_typed_statement ~stmt:(NRFunApp (StanLib, id, es)) ~return_type:NoReturnType ~loc + |> ok | Some (UnsizedType.ReturnType _) -> - Semantic_error.nonreturning_fn_expected_returning_found loc id.name - |> error + Semantic_error.nonreturning_fn_expected_returning_found loc id.name |> error | None -> - es - |> List.map ~f:type_of_expr_typed - |> Semantic_error.illtyped_stanlib_fn_app loc id.name - |> error) + es + |> List.map ~f:type_of_expr_typed + |> Semantic_error.illtyped_stanlib_fn_app loc id.name + |> error) +;; let semantic_check_nr_fnkind ~loc id es = match fn_kind_from_application id es with | StanLib -> semantic_check_nrfn_stan_math ~loc id es | UserDefined -> semantic_check_nrfn_normal ~loc id es +;; let semantic_check_nr_fn_app ~loc ~cf id es = Validate.( @@ -929,14 +1055,16 @@ let semantic_check_nr_fn_app ~loc ~cf id es = |> apply_const (semantic_check_identifier id) |> apply_const (semantic_check_nrfn_target ~loc ~cf id) >>= semantic_check_nr_fnkind ~loc id) +;; (* -- Assignment ------------------------------------------------------------ *) let semantic_check_assignment_read_only ~loc id = Validate.( - if Symbol_table.get_read_only vm id.name then - Semantic_error.cannot_assign_to_read_only loc id.name |> error + if Symbol_table.get_read_only vm id.name + then Semantic_error.cannot_assign_to_read_only loc id.name |> error else ok ()) +;; (* Variables from previous blocks are read-only. In particular, data and parameters never assigned to @@ -946,36 +1074,39 @@ let semantic_check_assignment_global ~loc ~cf ~block id = if (not (Symbol_table.is_global vm id.name)) || block = cf.current_block then ok () else Semantic_error.cannot_assign_to_global loc id.name |> error) +;; let mk_assignment_from_indexed_expr assop lhs rhs = - Assignment - {assign_lhs= Ast.lvalue_of_expr lhs; assign_op= assop; assign_rhs= rhs} + Assignment { assign_lhs = Ast.lvalue_of_expr lhs; assign_op = assop; assign_rhs = rhs } +;; let semantic_check_assignment_operator ~loc assop lhs rhs = Validate.( let err = - Semantic_error.illtyped_assignment loc assop lhs.emeta.type_ - rhs.emeta.type_ + Semantic_error.illtyped_assignment loc assop lhs.emeta.type_ rhs.emeta.type_ in match assop with | Assign | ArrowAssign -> - if - UnsizedType.check_of_same_type_mod_array_conv "" lhs.emeta.type_ - rhs.emeta.type_ - then - mk_typed_statement ~return_type:NoReturnType ~loc - ~stmt:(mk_assignment_from_indexed_expr assop lhs rhs) - |> ok - else error err + if UnsizedType.check_of_same_type_mod_array_conv "" lhs.emeta.type_ rhs.emeta.type_ + then + mk_typed_statement + ~return_type:NoReturnType + ~loc + ~stmt:(mk_assignment_from_indexed_expr assop lhs rhs) + |> ok + else error err | OperatorAssign op -> - List.map ~f:arg_type [lhs; rhs] - |> Stan_math_signatures.assignmentoperator_stan_math_return_type op - |> Option.value_map ~default:(error err) ~f:(function + List.map ~f:arg_type [ lhs; rhs ] + |> Stan_math_signatures.assignmentoperator_stan_math_return_type op + |> Option.value_map ~default:(error err) ~f:(function | ReturnType _ -> error err | Void -> - mk_typed_statement ~return_type:NoReturnType ~loc - ~stmt:(mk_assignment_from_indexed_expr assop lhs rhs) - |> ok )) + mk_typed_statement + ~return_type:NoReturnType + ~loc + ~stmt:(mk_assignment_from_indexed_expr assop lhs rhs) + |> ok)) +;; let semantic_check_assignment ~loc ~cf assign_lhs assign_op assign_rhs = let assign_id = Ast.id_of_lvalue assign_lhs in @@ -987,11 +1118,9 @@ let semantic_check_assignment ~loc ~cf assign_lhs assign_op assign_rhs = |> Option.map ~f:(fun (block, _) -> Validate.ok block) |> Option.value ~default: - ( if Stan_math_signatures.is_stan_math_function_name assign_id.name + (if Stan_math_signatures.is_stan_math_function_name assign_id.name then Validate.ok MathLibrary - else - Validate.error - @@ Semantic_error.ident_not_in_scope loc assign_id.name ) + else Validate.error @@ Semantic_error.ident_not_in_scope loc assign_id.name) in Validate.( liftA2 tuple2 (liftA3 tuple3 lhs assop rhs) block @@ -999,21 +1128,22 @@ let semantic_check_assignment ~loc ~cf assign_lhs assign_op assign_rhs = semantic_check_assignment_operator ~loc assop lhs rhs |> apply_const (semantic_check_assignment_global ~loc ~cf ~block assign_id) |> apply_const (semantic_check_assignment_read_only ~loc assign_id)) +;; (* -- Target plus-equals / Increment log-prob ------------------------------- *) let semantic_check_target_pe_expr_type ~loc e = match e.emeta.type_ with | UFun _ | UMathLibraryFunction -> - Semantic_error.int_or_real_container_expected loc e.emeta.type_ - |> Validate.error + Semantic_error.int_or_real_container_expected loc e.emeta.type_ |> Validate.error | _ -> Validate.ok () +;; let semantic_check_target_pe_usage ~loc ~cf = - if cf.in_lp_fun_def || cf.current_block = Model then Validate.ok () - else - Semantic_error.target_plusequals_outisde_model_or_logprob loc - |> Validate.error + if cf.in_lp_fun_def || cf.current_block = Model + then Validate.ok () + else Semantic_error.target_plusequals_outisde_model_or_logprob loc |> Validate.error +;; let semantic_check_target_pe ~loc ~cf e = Validate.( @@ -1022,8 +1152,8 @@ let semantic_check_target_pe ~loc ~cf e = >>= fun ue -> semantic_check_target_pe_expr_type ~loc ue |> map ~f:(fun _ -> - mk_typed_statement ~stmt:(TargetPE ue) ~return_type:NoReturnType - ~loc )) + mk_typed_statement ~stmt:(TargetPE ue) ~return_type:NoReturnType ~loc)) +;; let semantic_check_incr_logprob ~loc ~cf e = Validate.( @@ -1032,33 +1162,32 @@ let semantic_check_incr_logprob ~loc ~cf e = >>= fun ue -> semantic_check_target_pe_expr_type ~loc ue |> map ~f:(fun _ -> - mk_typed_statement ~stmt:(IncrementLogProb ue) - ~return_type:NoReturnType ~loc )) + mk_typed_statement ~stmt:(IncrementLogProb ue) ~return_type:NoReturnType ~loc)) +;; (* -- Tilde (Sampling notation) --------------------------------------------- *) let semantic_check_sampling_pdf_pmf id = Validate.( - if - String.( - is_suffix id.name ~suffix:"_lpdf" || is_suffix id.name ~suffix:"_lpmf") + if String.(is_suffix id.name ~suffix:"_lpdf" || is_suffix id.name ~suffix:"_lpmf") then error @@ Semantic_error.invalid_sampling_pdf_or_pmf id.id_loc else ok ()) +;; let semantic_check_sampling_cdf_ccdf ~loc id = Validate.( - if - String.( - is_suffix id.name ~suffix:"_cdf" || is_suffix id.name ~suffix:"_ccdf") + if String.(is_suffix id.name ~suffix:"_cdf" || is_suffix id.name ~suffix:"_ccdf") then error @@ Semantic_error.invalid_sampling_cdf_or_ccdf loc id.name else ok ()) +;; (* Target+= can only be used in model and functions with right suffix (same for tilde etc) *) let semantic_check_valid_sampling_pos ~loc ~cf = Validate.( - if not (cf.in_lp_fun_def || cf.current_block = Model) then - error @@ Semantic_error.target_plusequals_outisde_model_or_logprob loc + if not (cf.in_lp_fun_def || cf.current_block = Model) + then error @@ Semantic_error.target_plusequals_outisde_model_or_logprob loc else ok ()) +;; let semantic_check_sampling_distribution ~loc id arguments = let name = id.name @@ -1073,22 +1202,21 @@ let semantic_check_sampling_distribution ~loc id arguments = and valid_arg_types_for_suffix suffix = match Symbol_table.look vm (name ^ suffix) with | Some (Functions, UFun (listedtypes, ReturnType UReal)) -> - UnsizedType.check_compatible_arguments_mod_conv name listedtypes - argumenttypes + UnsizedType.check_compatible_arguments_mod_conv name listedtypes argumenttypes | _ -> false in Validate.( - if - is_reat_rt_for_suffix "_lpdf" - || is_reat_rt_for_suffix "_lpmf" - || is_reat_rt_for_suffix "_log" - && name <> "binomial_coefficient" - && name <> "multiply" - || valid_arg_types_for_suffix "_lpdf" - || valid_arg_types_for_suffix "_lpmf" - || valid_arg_types_for_suffix "_log" + if is_reat_rt_for_suffix "_lpdf" + || is_reat_rt_for_suffix "_lpmf" + || (is_reat_rt_for_suffix "_log" + && name <> "binomial_coefficient" + && name <> "multiply") + || valid_arg_types_for_suffix "_lpdf" + || valid_arg_types_for_suffix "_lpmf" + || valid_arg_types_for_suffix "_log" then ok () else error @@ Semantic_error.invalid_sampling_no_such_dist loc name) +;; let cumulative_density_is_defined id arguments = let name = id.name @@ -1103,38 +1231,38 @@ let cumulative_density_is_defined id arguments = and valid_arg_types_for_suffix suffix = match Symbol_table.look vm (name ^ suffix) with | Some (Functions, UFun (listedtypes, ReturnType UReal)) -> - UnsizedType.check_compatible_arguments_mod_conv name listedtypes - argumenttypes + UnsizedType.check_compatible_arguments_mod_conv name listedtypes argumenttypes | _ -> false in - ( is_reat_rt_for_suffix "_lcdf" + (is_reat_rt_for_suffix "_lcdf" || valid_arg_types_for_suffix "_lcdf" || is_reat_rt_for_suffix "_cdf_log" - || valid_arg_types_for_suffix "_cdf_log" ) - && ( is_reat_rt_for_suffix "_lccdf" + || valid_arg_types_for_suffix "_cdf_log") + && (is_reat_rt_for_suffix "_lccdf" || valid_arg_types_for_suffix "_lccdf" || is_reat_rt_for_suffix "_ccdf_log" - || valid_arg_types_for_suffix "_ccdf_log" ) + || valid_arg_types_for_suffix "_ccdf_log") +;; let can_truncate_distribution ~loc (arg : typed_expression) = function | NoTruncate -> Validate.ok () | _ -> - if UnsizedType.is_scalar_type arg.emeta.type_ then Validate.ok () - else Validate.error @@ Semantic_error.multivariate_truncation loc + if UnsizedType.is_scalar_type arg.emeta.type_ + then Validate.ok () + else Validate.error @@ Semantic_error.multivariate_truncation loc +;; let semantic_check_sampling_cdf_defined ~loc id truncation args = Validate.( match truncation with | NoTruncate -> ok () - | TruncateUpFrom e when cumulative_density_is_defined id (e :: args) -> - ok () - | TruncateDownFrom e when cumulative_density_is_defined id (e :: args) -> - ok () + | TruncateUpFrom e when cumulative_density_is_defined id (e :: args) -> ok () + | TruncateDownFrom e when cumulative_density_is_defined id (e :: args) -> ok () | TruncateBetween (e1, e2) when cumulative_density_is_defined id (e1 :: args) - && cumulative_density_is_defined id (e2 :: args) -> - ok () + && cumulative_density_is_defined id (e2 :: args) -> ok () | _ -> error @@ Semantic_error.invalid_truncation_cdf_or_ccdf loc) +;; let semantic_check_tilde ~loc ~cf distribution truncation arg args = Validate.( @@ -1148,27 +1276,31 @@ let semantic_check_tilde ~loc ~cf distribution truncation arg args = |> apply_const (semantic_check_sampling_cdf_ccdf ~loc distribution) >>= fun (truncation, arg, args) -> semantic_check_sampling_distribution ~loc distribution (arg :: args) - |> apply_const - (semantic_check_sampling_cdf_defined ~loc distribution truncation args) + |> apply_const (semantic_check_sampling_cdf_defined ~loc distribution truncation args) |> apply_const (can_truncate_distribution ~loc arg truncation) |> map ~f:(fun _ -> - let stmt = Tilde {arg; distribution; args; truncation} in - mk_typed_statement ~stmt ~loc ~return_type:NoReturnType )) + let stmt = Tilde { arg; distribution; args; truncation } in + mk_typed_statement ~stmt ~loc ~return_type:NoReturnType)) +;; (* -- Break ----------------------------------------------------------------- *) (* Break and continue only occur in loops. *) let semantic_check_break ~loc ~cf = Validate.( - if cf.loop_depth = 0 then Semantic_error.break_outside_loop loc |> error + if cf.loop_depth = 0 + then Semantic_error.break_outside_loop loc |> error else mk_typed_statement ~stmt:Break ~return_type:NoReturnType ~loc |> ok) +;; (* -- Continue -------------------------------------------------------------- *) let semantic_check_continue ~loc ~cf = Validate.( (* Break and continue only occur in loops. *) - if cf.loop_depth = 0 then Semantic_error.continue_outside_loop loc |> error + if cf.loop_depth = 0 + then Semantic_error.continue_outside_loop loc |> error else mk_typed_statement ~stmt:Continue ~return_type:NoReturnType ~loc |> ok) +;; (* -- Return ---------------------------------------------------------------- *) @@ -1177,23 +1309,25 @@ let semantic_check_continue ~loc ~cf = *) let semantic_check_return ~loc ~cf e = Validate.( - if not cf.in_returning_fun_def then - Semantic_error.expression_return_outside_returning_fn loc |> error + if not cf.in_returning_fun_def + then Semantic_error.expression_return_outside_returning_fn loc |> error else semantic_check_expression cf e |> map ~f:(fun ue -> - mk_typed_statement ~stmt:(Return ue) - ~return_type:(Complete (ReturnType ue.emeta.type_)) ~loc )) + mk_typed_statement + ~stmt:(Return ue) + ~return_type:(Complete (ReturnType ue.emeta.type_)) + ~loc)) +;; (* -- Return `void` --------------------------------------------------------- *) let semantic_check_returnvoid ~loc ~cf = Validate.( - if (not cf.in_fun_def) || cf.in_returning_fun_def then - Semantic_error.void_ouside_nonreturning_fn loc |> error - else - mk_typed_statement ~stmt:ReturnVoid ~return_type:(Complete Void) ~loc - |> ok) + if (not cf.in_fun_def) || cf.in_returning_fun_def + then Semantic_error.void_ouside_nonreturning_fn loc |> error + else mk_typed_statement ~stmt:ReturnVoid ~return_type:(Complete Void) ~loc |> ok) +;; (* -- Print ----------------------------------------------------------------- *) @@ -1203,8 +1337,8 @@ let semantic_check_print ~loc ~cf ps = |> List.map ~f:(semantic_check_printable cf) |> sequence |> map ~f:(fun ups -> - mk_typed_statement ~stmt:(Print ups) ~return_type:NoReturnType ~loc - )) + mk_typed_statement ~stmt:(Print ups) ~return_type:NoReturnType ~loc)) +;; (* -- Reject ---------------------------------------------------------------- *) @@ -1214,49 +1348,46 @@ let semantic_check_reject ~loc ~cf ps = |> List.map ~f:(semantic_check_printable cf) |> sequence |> map ~f:(fun ups -> - mk_typed_statement ~stmt:(Reject ups) ~return_type:AnyReturnType - ~loc )) + mk_typed_statement ~stmt:(Reject ups) ~return_type:AnyReturnType ~loc)) +;; (* -- Skip ------------------------------------------------------------------ *) let semantic_check_skip ~loc = mk_typed_statement ~stmt:Skip ~return_type:NoReturnType ~loc |> Validate.ok +;; (* -- If-Then-Else ---------------------------------------------------------- *) let try_compute_ifthenelse_statement_returntype loc srt1 srt2 = - match (srt1, srt2) with + match srt1, srt2 with | Complete rt1, Complete rt2 -> - lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Complete t) + lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Complete t) | Incomplete rt1, Incomplete rt2 - |Complete rt1, Incomplete rt2 - |Incomplete rt1, Complete rt2 -> - lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Incomplete t) - | AnyReturnType, NoReturnType - |NoReturnType, AnyReturnType - |NoReturnType, NoReturnType -> - Validate.ok NoReturnType + | Complete rt1, Incomplete rt2 + | Incomplete rt1, Complete rt2 -> + lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Incomplete t) + | AnyReturnType, NoReturnType | NoReturnType, AnyReturnType | NoReturnType, NoReturnType + -> Validate.ok NoReturnType | AnyReturnType, Incomplete rt - |Incomplete rt, AnyReturnType - |Complete rt, NoReturnType - |NoReturnType, Complete rt - |NoReturnType, Incomplete rt - |Incomplete rt, NoReturnType -> - Validate.ok @@ Incomplete rt - | Complete rt, AnyReturnType | AnyReturnType, Complete rt -> - Validate.ok @@ Complete rt + | Incomplete rt, AnyReturnType + | Complete rt, NoReturnType + | NoReturnType, Complete rt + | NoReturnType, Incomplete rt + | Incomplete rt, NoReturnType -> Validate.ok @@ Incomplete rt + | Complete rt, AnyReturnType | AnyReturnType, Complete rt -> Validate.ok @@ Complete rt | AnyReturnType, AnyReturnType -> Validate.ok AnyReturnType +;; let rec semantic_check_if_then_else ~loc ~cf pred_e s_true s_false_opt = let us1 = semantic_check_statement cf s_true and uos2 = s_false_opt |> Option.map ~f:(fun s -> - semantic_check_statement cf s |> Validate.map ~f:Option.some ) + semantic_check_statement cf s |> Validate.map ~f:Option.some) |> Option.value ~default:(Validate.ok None) and ue = - semantic_check_expression_of_int_or_real_type cf pred_e - "Condition in conditional" + semantic_check_expression_of_int_or_real_type cf pred_e "Condition in conditional" in Validate.( liftA3 tuple3 ue us1 uos2 @@ -1273,40 +1404,34 @@ let rec semantic_check_if_then_else ~loc ~cf pred_e s_true s_false_opt = (* -- While Statements ------------------------------------------------------ *) and semantic_check_while ~loc ~cf e s = - let us = semantic_check_statement {cf with loop_depth= cf.loop_depth + 1} s - and ue = - semantic_check_expression_of_int_or_real_type cf e - "Condition in while-loop" - in + let us = semantic_check_statement { cf with loop_depth = cf.loop_depth + 1 } s + and ue = semantic_check_expression_of_int_or_real_type cf e "Condition in while-loop" in Validate.liftA2 (fun ue us -> - mk_typed_statement - ~stmt:(While (ue, us)) - ~return_type:us.smeta.return_type ~loc ) - ue us + mk_typed_statement ~stmt:(While (ue, us)) ~return_type:us.smeta.return_type ~loc) + ue + us (* -- For Statements -------------------------------------------------------- *) and semantic_check_loop_body ~cf loop_var loop_var_ty loop_body = - Symbol_table.begin_scope vm ; + Symbol_table.begin_scope vm; let is_fresh_var = check_fresh_variable loop_var false in - Symbol_table.enter vm loop_var.name (cf.current_block, loop_var_ty) ; + Symbol_table.enter vm loop_var.name (cf.current_block, loop_var_ty); (* Check that function args and loop identifiers are not modified in function. (passed by const ref) *) - Symbol_table.set_read_only vm loop_var.name ; + Symbol_table.set_read_only vm loop_var.name; let us = - semantic_check_statement {cf with loop_depth= cf.loop_depth + 1} loop_body + semantic_check_statement { cf with loop_depth = cf.loop_depth + 1 } loop_body |> Validate.apply_const is_fresh_var in - Symbol_table.end_scope vm ; us + Symbol_table.end_scope vm; + us -and semantic_check_for ~loc ~cf loop_var lower_bound_e upper_bound_e loop_body - = +and semantic_check_for ~loc ~cf loop_var lower_bound_e upper_bound_e loop_body = let ue1 = - semantic_check_expression_of_int_type cf lower_bound_e - "Lower bound of for-loop" + semantic_check_expression_of_int_type cf lower_bound_e "Lower bound of for-loop" and ue2 = - semantic_check_expression_of_int_type cf upper_bound_e - "Upper bound of for-loop" + semantic_check_expression_of_int_type cf upper_bound_e "Upper bound of for-loop" in Validate.( liftA2 tuple2 ue1 ue2 @@ -1317,11 +1442,13 @@ and semantic_check_for ~loc ~cf loop_var lower_bound_e upper_bound_e loop_body mk_typed_statement ~stmt: (For - { loop_variable= loop_var - ; lower_bound= ue1 - ; upper_bound= ue2 - ; loop_body= us }) - ~return_type:us.smeta.return_type ~loc )) + { loop_variable = loop_var + ; lower_bound = ue1 + ; upper_bound = ue2 + ; loop_body = us + }) + ~return_type:us.smeta.return_type + ~loc)) (* -- Foreach Statements ---------------------------------------------------- *) and semantic_check_foreach_loop_identifier_type ~loc ty = @@ -1329,85 +1456,80 @@ and semantic_check_foreach_loop_identifier_type ~loc ty = match ty with | UnsizedType.UArray ut -> ok ut | UVector | URowVector | UMatrix -> ok UnsizedType.UReal - | _ -> - Semantic_error.array_vector_rowvector_matrix_expected loc ty |> error) + | _ -> Semantic_error.array_vector_rowvector_matrix_expected loc ty |> error) and semantic_check_foreach ~loc ~cf loop_var foreach_expr loop_body = Validate.( semantic_check_expression cf foreach_expr |> apply_const (semantic_check_identifier loop_var) >>= fun ue -> - semantic_check_foreach_loop_identifier_type ~loc:ue.emeta.loc - ue.emeta.type_ + semantic_check_foreach_loop_identifier_type ~loc:ue.emeta.loc ue.emeta.type_ >>= fun loop_var_ty -> semantic_check_loop_body ~cf loop_var loop_var_ty loop_body |> map ~f:(fun us -> mk_typed_statement ~stmt:(ForEach (loop_var, ue, us)) - ~return_type:us.smeta.return_type ~loc )) + ~return_type:us.smeta.return_type + ~loc)) (* -- Blocks ---------------------------------------------------------------- *) -and stmt_is_escape {stmt; _} = +and stmt_is_escape { stmt; _ } = match stmt with | Break | Continue | Reject _ | Return _ | ReturnVoid -> true | _ -> false and list_until_escape xs = let rec aux accu = function - | next :: next' :: _ when stmt_is_escape next' -> - List.rev (next' :: next :: accu) + | next :: next' :: _ when stmt_is_escape next' -> List.rev (next' :: next :: accu) | next :: rest -> aux (next :: accu) rest | [] -> List.rev accu in aux [] xs and try_compute_block_statement_returntype loc srt1 srt2 = - match (srt1, srt2) with + match srt1, srt2 with | Complete rt1, Complete rt2 | Incomplete rt1, Complete rt2 -> - lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Complete t) + lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Complete t) | Incomplete rt1, Incomplete rt2 | Complete rt1, Incomplete rt2 -> - lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Incomplete t) + lub_rt loc rt1 rt2 |> Validate.map ~f:(fun t -> Incomplete t) | NoReturnType, NoReturnType -> Validate.ok NoReturnType | AnyReturnType, Incomplete rt - |Complete rt, NoReturnType - |NoReturnType, Incomplete rt - |Incomplete rt, NoReturnType -> - Validate.ok @@ Incomplete rt + | Complete rt, NoReturnType + | NoReturnType, Incomplete rt + | Incomplete rt, NoReturnType -> Validate.ok @@ Incomplete rt | NoReturnType, Complete rt - |Complete rt, AnyReturnType - |Incomplete rt, AnyReturnType - |AnyReturnType, Complete rt -> - Validate.ok @@ Complete rt + | Complete rt, AnyReturnType + | Incomplete rt, AnyReturnType + | AnyReturnType, Complete rt -> Validate.ok @@ Complete rt | AnyReturnType, NoReturnType - |NoReturnType, AnyReturnType - |AnyReturnType, AnyReturnType -> - Validate.ok AnyReturnType + | NoReturnType, AnyReturnType + | AnyReturnType, AnyReturnType -> Validate.ok AnyReturnType and semantic_check_block ~loc ~cf stmts = - Symbol_table.begin_scope vm ; + Symbol_table.begin_scope vm; (* Any statements after a break or continue or return or reject do not count for the return type. *) let validated_stmts = List.map ~f:(semantic_check_statement cf) stmts |> Validate.sequence in - Symbol_table.end_scope vm ; + Symbol_table.end_scope vm; Validate.( validated_stmts >>= fun xs -> let return_ty = - xs |> list_until_escape + xs + |> list_until_escape |> List.map ~f:(fun s -> s.smeta.return_type) |> List.fold ~init:(ok NoReturnType) ~f:(fun accu x -> - accu >>= fun y -> try_compute_block_statement_returntype loc y x - ) + accu >>= fun y -> try_compute_block_statement_returntype loc y x) in map return_ty ~f:(fun return_type -> - mk_typed_statement ~stmt:(Block xs) ~return_type ~loc )) + mk_typed_statement ~stmt:(Block xs) ~return_type ~loc)) (* -- Variable Declarations ------------------------------------------------- *) and semantic_check_var_decl_bounds ~loc is_global sized_ty trans = - let is_real {emeta; _} = emeta.type_ = UReal in + let is_real { emeta; _ } = emeta.type_ = UReal in let is_valid_transformation = match trans with | Program.Lower e -> is_real e @@ -1416,16 +1538,15 @@ and semantic_check_var_decl_bounds ~loc is_global sized_ty trans = | _ -> false in Validate.( - if is_global && sized_ty = SizedType.SInt && is_valid_transformation then - Semantic_error.non_int_bounds loc |> error + if is_global && sized_ty = SizedType.SInt && is_valid_transformation + then Semantic_error.non_int_bounds loc |> error else ok ()) and semantic_check_transformed_param_ty ~loc ~cf is_global unsized_ty = Validate.( - if - is_global - && (cf.current_block = Param || cf.current_block = TParam) - && UnsizedType.contains_int unsized_ty + if is_global + && (cf.current_block = Param || cf.current_block = TParam) + && UnsizedType.contains_int unsized_ty then Semantic_error.transformed_params_int loc |> error else ok ()) @@ -1434,92 +1555,92 @@ and semantic_check_var_decl_initial_value ~loc ~cf id init_val_opt = |> Option.value_map ~default:(Validate.ok None) ~f:(fun e -> let stmt = Assignment - { assign_lhs= {lval= LVariable id; lmeta= {loc}} - ; assign_op= Assign - ; assign_rhs= e } + { assign_lhs = { lval = LVariable id; lmeta = { loc } } + ; assign_op = Assign + ; assign_rhs = e + } in mk_untyped_statement ~loc ~stmt |> semantic_check_statement cf |> Validate.map ~f:(fun ts -> - match (ts.stmt, ts.smeta.return_type) with - | Assignment {assign_rhs= ue; _}, NoReturnType -> Some ue + match ts.stmt, ts.smeta.return_type with + | Assignment { assign_rhs = ue; _ }, NoReturnType -> Some ue | _ -> - let msg = - "semantic_check_var_decl: `Assignment` expected." - in - fatal_error ~msg () ) ) + let msg = "semantic_check_var_decl: `Assignment` expected." in + fatal_error ~msg ())) and semantic_check_var_decl ~loc ~cf sized_ty trans id init is_global = let checked_stmt = - semantic_check_sizedtype {cf with in_toplevel_decl= is_global} sized_ty + semantic_check_sizedtype { cf with in_toplevel_decl = is_global } sized_ty in Validate.( let checked_trans = checked_stmt - >>= fun ust -> - semantic_check_transformation cf (SizedType.to_unsized ust) trans + >>= fun ust -> semantic_check_transformation cf (SizedType.to_unsized ust) trans in liftA2 tuple2 checked_stmt checked_trans |> apply_const (semantic_check_identifier id) |> apply_const (check_fresh_variable id false) >>= fun (ust, utrans) -> let ut = SizedType.to_unsized ust in - Symbol_table.enter vm id.name (cf.current_block, ut) ; + Symbol_table.enter vm id.name (cf.current_block, ut); semantic_check_var_decl_initial_value ~loc ~cf id init |> apply_const (semantic_check_var_decl_bounds ~loc is_global ust utrans) |> apply_const (semantic_check_transformed_param_ty ~loc ~cf is_global ut) |> map ~f:(fun uinit -> let stmt = VarDecl - { decl_type= Sized ust - ; transformation= utrans - ; identifier= id - ; initial_value= uinit - ; is_global } + { decl_type = Sized ust + ; transformation = utrans + ; identifier = id + ; initial_value = uinit + ; is_global + } in - mk_typed_statement ~stmt ~loc ~return_type:NoReturnType )) + mk_typed_statement ~stmt ~loc ~return_type:NoReturnType)) (* -- Function definitions -------------------------------------------------- *) and semantic_check_fundef_overloaded ~loc id arg_tys rt = Validate.( (* User defined functions cannot be overloaded *) - if Symbol_table.check_is_unassigned vm id.name then + if Symbol_table.check_is_unassigned vm id.name + then ( match Symbol_table.look vm id.name with - | Some (Functions, UFun (arg_tys', rt')) - when arg_tys' = arg_tys && rt' = rt -> - ok () + | Some (Functions, UFun (arg_tys', rt')) when arg_tys' = arg_tys && rt' = rt -> + ok () | _ -> - Symbol_table.look vm id.name - |> Option.map ~f:snd - |> Semantic_error.mismatched_fn_def_decl loc id.name - |> error + Symbol_table.look vm id.name + |> Option.map ~f:snd + |> Semantic_error.mismatched_fn_def_decl loc id.name + |> error) else check_fresh_variable id (List.length arg_tys = 0)) (** WARNING: side effecting *) and semantic_check_fundef_decl ~loc id body = Validate.( match body with - | {stmt= Skip; _} -> - if Symbol_table.check_is_unassigned vm id.name then - error @@ Semantic_error.fn_decl_without_def loc - else - let () = Symbol_table.set_is_unassigned vm id.name in - ok () - | _ -> - Symbol_table.set_is_assigned vm id.name ; + | { stmt = Skip; _ } -> + if Symbol_table.check_is_unassigned vm id.name + then error @@ Semantic_error.fn_decl_without_def loc + else ( + let () = Symbol_table.set_is_unassigned vm id.name in ok ()) + | _ -> + Symbol_table.set_is_assigned vm id.name; + ok ()) and semantic_check_fundef_dist_rt ~loc id return_ty = Validate.( let is_dist = List.exists ~f:(fun x -> String.is_suffix id.name ~suffix:x) - ["_log"; "_lpdf"; "_lpmf"; "_lcdf"; "_lccdf"] + [ "_log"; "_lpdf"; "_lpmf"; "_lcdf"; "_lccdf" ] in - if is_dist then + if is_dist + then ( match return_ty with | UnsizedType.ReturnType UReal -> ok () - | _ -> error @@ Semantic_error.non_real_prob_fn_def loc + | _ -> error @@ Semantic_error.non_real_prob_fn_def loc) else ok ()) and semantic_check_pdf_fundef_first_arg_ty ~loc id arg_tys = @@ -1527,23 +1648,21 @@ and semantic_check_pdf_fundef_first_arg_ty ~loc id arg_tys = (* TODO: I think these kind of functions belong with the type definition *) let is_real_type = function | UnsizedType.UReal | UVector | URowVector | UMatrix - |UArray UReal - |UArray UVector - |UArray URowVector - |UArray UMatrix -> - true + | UArray UReal + | UArray UVector + | UArray URowVector + | UArray UMatrix -> true | _ -> false in - if String.is_suffix id.name ~suffix:"_lpdf" then + if String.is_suffix id.name ~suffix:"_lpdf" + then List.hd arg_tys |> Option.value_map - ~default: - (error @@ Semantic_error.prob_density_non_real_variate loc None) + ~default:(error @@ Semantic_error.prob_density_non_real_variate loc None) ~f:(fun (_, rt) -> - if is_real_type rt then ok () - else - error - @@ Semantic_error.prob_density_non_real_variate loc (Some rt) ) + if is_real_type rt + then ok () + else error @@ Semantic_error.prob_density_non_real_variate loc (Some rt)) else ok ()) and semantic_check_pmf_fundef_first_arg_ty ~loc id arg_tys = @@ -1553,30 +1672,29 @@ and semantic_check_pmf_fundef_first_arg_ty ~loc id arg_tys = | UnsizedType.UInt | UArray UInt -> true | _ -> false in - if String.is_suffix id.name ~suffix:"_lpmf" then + if String.is_suffix id.name ~suffix:"_lpmf" + then List.hd arg_tys |> Option.value_map ~default:(error @@ Semantic_error.prob_mass_non_int_variate loc None) ~f:(fun (_, rt) -> - if is_int_type rt then ok () - else - error @@ Semantic_error.prob_mass_non_int_variate loc (Some rt) - ) + if is_int_type rt + then ok () + else error @@ Semantic_error.prob_mass_non_int_variate loc (Some rt)) else ok ()) (* All function arguments are distinct *) and semantic_check_fundef_distinct_arg_ids ~loc arg_names = Validate.( - if dup_exists arg_names then - error @@ Semantic_error.duplicate_arg_names loc + if dup_exists arg_names + then error @@ Semantic_error.duplicate_arg_names loc else ok ()) (* Check that every trace through function body contains return statement of right type *) and semantic_check_fundef_return_tys ~loc id return_type body = Validate.( - if - Symbol_table.check_is_unassigned vm id.name - || check_of_compatible_return_type return_type body.smeta.return_type + if Symbol_table.check_is_unassigned vm id.name + || check_of_compatible_return_type return_type body.smeta.return_type then ok () else error @@ Semantic_error.incompatible_return_types loc) @@ -1587,7 +1705,7 @@ and semantic_check_fundef ~loc ~cf return_ty id args body = semantic_check_autodifftype at |> apply_const (semantic_check_unsizedtype ut) |> apply_const (semantic_check_identifier id) - |> map ~f:(fun at -> (at, ut, id))) ) + |> map ~f:(fun at -> at, ut, id))) |> Validate.sequence in Validate.( @@ -1596,23 +1714,23 @@ and semantic_check_fundef ~loc ~cf return_ty id args body = |> apply_const (semantic_check_returntype return_ty) >>= fun uargs -> let urt = return_ty in - let uarg_types = List.map ~f:(fun (w, y, _) -> (w, y)) uargs in + let uarg_types = List.map ~f:(fun (w, y, _) -> w, y) uargs in let uarg_identifiers = List.map ~f:(fun (_, _, z) -> z) uargs in let uarg_names = List.map ~f:(fun x -> x.name) uarg_identifiers in semantic_check_fundef_overloaded ~loc id uarg_types urt |> apply_const (semantic_check_fundef_decl ~loc id body) >>= fun _ -> (* WARNING: SIDE EFFECTING *) - Symbol_table.enter vm id.name (Functions, UFun (uarg_types, urt)) ; + Symbol_table.enter vm id.name (Functions, UFun (uarg_types, urt)); (* Check that function args and loop identifiers are not modified in function. (passed by const ref)*) - List.iter ~f:(Symbol_table.set_read_only vm) uarg_names ; + List.iter ~f:(Symbol_table.set_read_only vm) uarg_names; semantic_check_fundef_dist_rt ~loc id urt |> apply_const (semantic_check_pdf_fundef_first_arg_ty ~loc id uarg_types) |> apply_const (semantic_check_pmf_fundef_first_arg_ty ~loc id uarg_types) >>= fun _ -> (* WARNING: SIDE EFFECTING *) - Symbol_table.begin_scope vm ; + Symbol_table.begin_scope vm; List.map ~f:(fun x -> check_fresh_variable x false) uarg_identifiers |> sequence |> apply_const (semantic_check_fundef_distinct_arg_ids ~loc uarg_names) @@ -1625,19 +1743,22 @@ and semantic_check_fundef ~loc ~cf return_ty id args body = as if they are parameters, for the purposes of type checking. *) (* WARNING: SIDE EFFECTING *) - let _ : unit Base.List.Or_unequal_lengths.t = - List.iter2 ~f:(Symbol_table.enter vm) uarg_names + let (_ : unit Base.List.Or_unequal_lengths.t) = + List.iter2 + ~f:(Symbol_table.enter vm) + uarg_names (List.map ~f:(function - | UnsizedType.DataOnly, ut -> (Data, ut) - | AutoDiffable, ut -> (Param, ut)) + | UnsizedType.DataOnly, ut -> Data, ut + | AutoDiffable, ut -> Param, ut) uarg_types) and context = { cf with - in_fun_def= true - ; in_rng_fun_def= String.is_suffix id.name ~suffix:"_rng" - ; in_lp_fun_def= String.is_suffix id.name ~suffix:"_lp" - ; in_returning_fun_def= urt <> Void } + in_fun_def = true + ; in_rng_fun_def = String.is_suffix id.name ~suffix:"_rng" + ; in_lp_fun_def = String.is_suffix id.name ~suffix:"_lp" + ; in_returning_fun_def = urt <> Void + } in let body' = semantic_check_statement context body in body' @@ -1645,24 +1766,25 @@ and semantic_check_fundef ~loc ~cf return_ty id args body = semantic_check_fundef_return_tys ~loc id urt ub |> map ~f:(fun _ -> (* WARNING: SIDE EFFECTING *) - Symbol_table.end_scope vm ; + Symbol_table.end_scope vm; let stmt = - FunDef {returntype= urt; funname= id; arguments= uargs; body= ub} + FunDef { returntype = urt; funname = id; arguments = uargs; body = ub } in - mk_typed_statement ~return_type:NoReturnType ~loc ~stmt )) + mk_typed_statement ~return_type:NoReturnType ~loc ~stmt)) (* -- Top-level Statements -------------------------------------------------- *) -and semantic_check_statement cf (s : Ast.untyped_statement) : - Ast.typed_statement Validate.t = +and semantic_check_statement cf (s : Ast.untyped_statement) + : Ast.typed_statement Validate.t + = let loc = s.smeta.loc in match s.stmt with | NRFunApp (_, id, es) -> semantic_check_nr_fn_app ~loc ~cf id es - | Assignment {assign_lhs; assign_op; assign_rhs} -> - semantic_check_assignment ~loc ~cf assign_lhs assign_op assign_rhs + | Assignment { assign_lhs; assign_op; assign_rhs } -> + semantic_check_assignment ~loc ~cf assign_lhs assign_op assign_rhs | TargetPE e -> semantic_check_target_pe ~loc ~cf e | IncrementLogProb e -> semantic_check_incr_logprob ~loc ~cf e - | Tilde {arg; distribution; args; truncation} -> - semantic_check_tilde ~loc ~cf distribution truncation arg args + | Tilde { arg; distribution; args; truncation } -> + semantic_check_tilde ~loc ~cf distribution truncation arg args | Break -> semantic_check_break ~loc ~cf | Continue -> semantic_check_continue ~loc ~cf | Return e -> semantic_check_return ~loc ~cf e @@ -1672,123 +1794,120 @@ and semantic_check_statement cf (s : Ast.untyped_statement) : | Skip -> semantic_check_skip ~loc | IfThenElse (e, s1, os2) -> semantic_check_if_then_else ~loc ~cf e s1 os2 | While (e, s) -> semantic_check_while ~loc ~cf e s - | For {loop_variable; lower_bound; upper_bound; loop_body} -> - semantic_check_for ~loc ~cf loop_variable lower_bound upper_bound - loop_body + | For { loop_variable; lower_bound; upper_bound; loop_body } -> + semantic_check_for ~loc ~cf loop_variable lower_bound upper_bound loop_body | ForEach (id, e, s) -> semantic_check_foreach ~loc ~cf id e s | Block vdsl -> semantic_check_block ~loc ~cf vdsl - | VarDecl {decl_type= Unsized _; _} -> - raise_s [%message "Don't support unsized declarations yet."] - | VarDecl - { decl_type= Sized st - ; transformation - ; identifier - ; initial_value - ; is_global } -> - semantic_check_var_decl ~loc ~cf st transformation identifier - initial_value is_global - | FunDef {returntype; funname; arguments; body} -> - semantic_check_fundef ~loc ~cf returntype funname arguments body + | VarDecl { decl_type = Unsized _; _ } -> + raise_s [%message "Don't support unsized declarations yet."] + | VarDecl { decl_type = Sized st; transformation; identifier; initial_value; is_global } + -> + semantic_check_var_decl ~loc ~cf st transformation identifier initial_value is_global + | FunDef { returntype; funname; arguments; body } -> + semantic_check_fundef ~loc ~cf returntype funname arguments body +;; (* == Untyped programs ====================================================== *) let semantic_check_ostatements_in_block ~cf block stmts_opt = - let cf' = {cf with current_block= block} in + let cf' = { cf with current_block = block } in Option.value_map stmts_opt ~default:(Validate.ok None) ~f:(fun stmts -> (* I'm folding since I'm not sure if map is guaranteed to respect the ordering of the list *) List.fold ~init:[] stmts ~f:(fun accu stmt -> let s = semantic_check_statement cf' stmt in - s :: accu ) - |> List.rev |> Validate.sequence - |> Validate.map ~f:Option.some ) + s :: accu) + |> List.rev + |> Validate.sequence + |> Validate.map ~f:Option.some) +;; let check_fun_def_body_in_block = function - | {stmt= FunDef {body= {stmt= Block _; _}; _}; _} - |{stmt= FunDef {body= {stmt= Skip; _}; _}; _} -> - Validate.ok () - | {stmt= FunDef {body= {stmt= _; smeta}; _}; _} -> - Validate.error @@ Semantic_error.fn_decl_needs_block smeta.loc + | { stmt = FunDef { body = { stmt = Block _; _ }; _ }; _ } + | { stmt = FunDef { body = { stmt = Skip; _ }; _ }; _ } -> Validate.ok () + | { stmt = FunDef { body = { stmt = _; smeta }; _ }; _ } -> + Validate.error @@ Semantic_error.fn_decl_needs_block smeta.loc | _ -> Validate.ok () +;; let semantic_check_functions_have_defn function_block_stmts_opt = Validate.( - if - Symbol_table.check_some_id_is_unassigned vm - && !check_that_all_functions_have_definition - then + if Symbol_table.check_some_id_is_unassigned vm + && !check_that_all_functions_have_definition + then ( match function_block_stmts_opt with - | Some ({smeta; _} :: _) -> - (* TODO: insert better location in the error *) - error @@ Semantic_error.fn_decl_without_def smeta.loc - | _ -> fatal_error ~msg:"semantic_check_functions_have_defn" () - else + | Some ({ smeta; _ } :: _) -> + (* TODO: insert better location in the error *) + error @@ Semantic_error.fn_decl_without_def smeta.loc + | _ -> fatal_error ~msg:"semantic_check_functions_have_defn" ()) + else ( match function_block_stmts_opt with | Some [] | None -> ok () | Some ls -> - List.map ~f:check_fun_def_body_in_block ls - |> sequence - |> map ~f:(fun _ -> ())) + List.map ~f:check_fun_def_body_in_block ls |> sequence |> map ~f:(fun _ -> ()))) +;; (* The actual semantic checks for all AST nodes! *) let semantic_check_program - { functionblock= fb - ; datablock= db - ; transformeddatablock= tdb - ; parametersblock= pb - ; transformedparametersblock= tpb - ; modelblock= mb - ; generatedquantitiesblock= gb } = + { functionblock = fb + ; datablock = db + ; transformeddatablock = tdb + ; parametersblock = pb + ; transformedparametersblock = tpb + ; modelblock = mb + ; generatedquantitiesblock = gb + } + = (* NB: We always want to make sure we start with an empty symbol table, in case we are processing multiple files in one run. *) - unsafe_clear_symbol_table vm ; + unsafe_clear_symbol_table vm; let cf = - { current_block= Functions - ; in_toplevel_decl= false - ; in_fun_def= false - ; in_returning_fun_def= false - ; in_rng_fun_def= false - ; in_lp_fun_def= false - ; loop_depth= 0 } + { current_block = Functions + ; in_toplevel_decl = false + ; in_fun_def = false + ; in_returning_fun_def = false + ; in_rng_fun_def = false + ; in_lp_fun_def = false + ; loop_depth = 0 + } in let ufb = Validate.( semantic_check_ostatements_in_block ~cf Functions fb - >>= fun xs -> - semantic_check_functions_have_defn xs |> map ~f:(fun _ -> xs)) + >>= fun xs -> semantic_check_functions_have_defn xs |> map ~f:(fun _ -> xs)) in let udb = semantic_check_ostatements_in_block ~cf Data db in let utdb = semantic_check_ostatements_in_block ~cf TData tdb in let upb = semantic_check_ostatements_in_block ~cf Param pb in let utpb = semantic_check_ostatements_in_block ~cf TParam tpb in (* Model top level variables only assigned and read in model *) - Symbol_table.begin_scope vm ; + Symbol_table.begin_scope vm; let umb = semantic_check_ostatements_in_block ~cf Model mb in - Symbol_table.end_scope vm ; + Symbol_table.end_scope vm; let ugb = semantic_check_ostatements_in_block ~cf GQuant gb in let mk_typed_prog ufb udb utdb upb utpb umb ugb : Ast.typed_program = - { functionblock= ufb - ; datablock= udb - ; transformeddatablock= utdb - ; parametersblock= upb - ; transformedparametersblock= utpb - ; modelblock= umb - ; generatedquantitiesblock= ugb } + { functionblock = ufb + ; datablock = udb + ; transformeddatablock = utdb + ; parametersblock = upb + ; transformedparametersblock = utpb + ; modelblock = umb + ; generatedquantitiesblock = ugb + } in let apply_to x f = Validate.apply ~f x in - let check_correctness_invariant (decorated_ast : typed_program) : - typed_program = - if - compare_untyped_program - { functionblock= fb - ; datablock= db - ; transformeddatablock= tdb - ; parametersblock= pb - ; transformedparametersblock= tpb - ; modelblock= mb - ; generatedquantitiesblock= gb } - (untyped_program_of_typed_program decorated_ast) - = 0 + let check_correctness_invariant (decorated_ast : typed_program) : typed_program = + if compare_untyped_program + { functionblock = fb + ; datablock = db + ; transformeddatablock = tdb + ; parametersblock = pb + ; transformedparametersblock = tpb + ; modelblock = mb + ; generatedquantitiesblock = gb + } + (untyped_program_of_typed_program decorated_ast) + = 0 then decorated_ast else raise_s @@ -1800,9 +1919,16 @@ let semantic_check_program Validate.map ~f:check_correctness_invariant in Validate.( - ok mk_typed_prog |> apply_to ufb |> apply_to udb |> apply_to utdb - |> apply_to upb |> apply_to utpb |> apply_to umb |> apply_to ugb + ok mk_typed_prog + |> apply_to ufb + |> apply_to udb + |> apply_to utdb + |> apply_to upb + |> apply_to utpb + |> apply_to umb + |> apply_to ugb |> check_correctness_invariant_validate |> get_with ~with_ok:(fun ok -> Result.Ok ok) ~with_errors:(fun errs -> Result.Error errs)) +;; diff --git a/src/frontend/Semantic_check.mli b/src/frontend/Semantic_check.mli index d635f5497c..d53b664da4 100644 --- a/src/frontend/Semantic_check.mli +++ b/src/frontend/Semantic_check.mli @@ -2,26 +2,27 @@ open Core_kernel -val inferred_unsizedtype_of_indexed_exn : - loc:Middle.Location_span.t +(** Infers unsized type of an `Indexed` expression *) +val inferred_unsizedtype_of_indexed_exn + : loc:Middle.Location_span.t -> Middle.UnsizedType.t -> Ast.typed_expression Ast.index list -> Middle.UnsizedType.t -(** Infers unsized type of an `Indexed` expression *) -val semantic_check_binop_exn : - Middle.Location_span.t +val semantic_check_binop_exn + : Middle.Location_span.t -> Middle.Operator.t -> Ast.typed_expression * Ast.typed_expression -> Ast.typed_expression -val semantic_check_program : - Ast.untyped_program -> (Ast.typed_program, Semantic_error.t list) result (** Performs semantic check on AST and returns original AST embellished with type decorations *) +val semantic_check_program + : Ast.untyped_program + -> (Ast.typed_program, Semantic_error.t list) result -val check_that_all_functions_have_definition : bool ref (** A switch to determine whether we check that all functions have a definition *) +val check_that_all_functions_have_definition : bool ref -val model_name : string ref (** A reference to hold the model name. Relevant for checking variable clashes and used in code generation. *) +val model_name : string ref diff --git a/src/frontend/Semantic_error.ml b/src/frontend/Semantic_error.ml index 8e40df6ad9..de81add036 100644 --- a/src/frontend/Semantic_error.ml +++ b/src/frontend/Semantic_error.ml @@ -14,13 +14,10 @@ module TypeError = struct | IntIntArrayOrRangeExpected of UnsizedType.t | IntOrRealContainerExpected of UnsizedType.t | ArrayVectorRowVectorMatrixExpected of UnsizedType.t - | IllTypedAssignment of - Ast.assignmentoperator * UnsizedType.t * UnsizedType.t + | IllTypedAssignment of Ast.assignmentoperator * UnsizedType.t * UnsizedType.t | IllTypedTernaryIf of UnsizedType.t * UnsizedType.t * UnsizedType.t | IllTypedReduceSum of - string - * UnsizedType.t list - * (UnsizedType.autodifftype * UnsizedType.t) list + string * UnsizedType.t list * (UnsizedType.autodifftype * UnsizedType.t) list | IllTypedReduceSumGeneric of string * UnsizedType.t list | ReturningFnExpectedNonReturningFound of string | ReturningFnExpectedNonFnFound of string @@ -41,197 +38,263 @@ module TypeError = struct let pp ppf = function | MismatchedReturnTypes (rt1, rt2) -> - Fmt.pf ppf - "Branches of function definition need to have the same return type. \ - Instead, found return types %a and %a." - UnsizedType.pp_returntype rt1 UnsizedType.pp_returntype rt2 + Fmt.pf + ppf + "Branches of function definition need to have the same return type. Instead, \ + found return types %a and %a." + UnsizedType.pp_returntype + rt1 + UnsizedType.pp_returntype + rt2 | MismatchedArrayTypes (t1, t2) -> - Fmt.pf ppf - "Array expression must have entries of consistent type. Expected %a \ - but found %a." - UnsizedType.pp t1 UnsizedType.pp t2 + Fmt.pf + ppf + "Array expression must have entries of consistent type. Expected %a but found %a." + UnsizedType.pp + t1 + UnsizedType.pp + t2 | InvalidRowVectorTypes ty -> - Fmt.pf ppf - "Row_vector expression must have all int or real entries. Found \ - type %a." - UnsizedType.pp ty + Fmt.pf + ppf + "Row_vector expression must have all int or real entries. Found type %a." + UnsizedType.pp + ty | InvalidMatrixTypes ty -> - Fmt.pf ppf - "Matrix expression must have all row_vector entries. Found type %a." - UnsizedType.pp ty + Fmt.pf + ppf + "Matrix expression must have all row_vector entries. Found type %a." + UnsizedType.pp + ty | IntExpected (name, ut) -> - Fmt.pf ppf "%s must be of type int. Instead found type %a." name - UnsizedType.pp ut + Fmt.pf ppf "%s must be of type int. Instead found type %a." name UnsizedType.pp ut | IntOrRealExpected (name, ut) -> - Fmt.pf ppf "%s must be of type int or real. Instead found type %a." - name UnsizedType.pp ut + Fmt.pf + ppf + "%s must be of type int or real. Instead found type %a." + name + UnsizedType.pp + ut | TypeExpected (name, (UInt | UReal), ut) -> - Fmt.pf ppf "%s must be a scalar. Instead found type %a." name - UnsizedType.pp ut + Fmt.pf ppf "%s must be a scalar. Instead found type %a." name UnsizedType.pp ut | TypeExpected (name, et, ut) -> - Fmt.pf ppf "%s must be a scalar or of type %a. Instead found type %a." - name UnsizedType.pp et UnsizedType.pp ut + Fmt.pf + ppf + "%s must be a scalar or of type %a. Instead found type %a." + name + UnsizedType.pp + et + UnsizedType.pp + ut | IntOrRealContainerExpected ut -> - Fmt.pf ppf - "A (container of) real or int was expected. Instead found type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "A (container of) real or int was expected. Instead found type %a." + UnsizedType.pp + ut | IntIntArrayOrRangeExpected ut -> - Fmt.pf ppf - "Index must be of type int or int[] or must be a range. Instead \ - found type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "Index must be of type int or int[] or must be a range. Instead found type %a." + UnsizedType.pp + ut | ArrayVectorRowVectorMatrixExpected ut -> - Fmt.pf ppf - "Foreach-loop must be over array, vector, row_vector or matrix. \ - Instead found expression of type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "Foreach-loop must be over array, vector, row_vector or matrix. Instead found \ + expression of type %a." + UnsizedType.pp + ut | IllTypedAssignment ((OperatorAssign op as assignop), lt, rt) -> - Fmt.pf ppf - "@[Ill-typed arguments supplied to assignment operator %s: lhs \ - has type %a and rhs has type %a. Available signatures:@]%s" - (Pretty_printing.pretty_print_assignmentoperator assignop) - UnsizedType.pp lt UnsizedType.pp rt - ( Stan_math_signatures.pretty_print_math_lib_assignmentoperator_sigs - op - |> Option.value ~default:"no matching signatures" ) + Fmt.pf + ppf + "@[Ill-typed arguments supplied to assignment operator %s: lhs has type %a \ + and rhs has type %a. Available signatures:@]%s" + (Pretty_printing.pretty_print_assignmentoperator assignop) + UnsizedType.pp + lt + UnsizedType.pp + rt + (Stan_math_signatures.pretty_print_math_lib_assignmentoperator_sigs op + |> Option.value ~default:"no matching signatures") | IllTypedAssignment (assignop, lt, rt) -> - Fmt.pf ppf - "Ill-typed arguments supplied to assignment operator %s: lhs has \ - type %a and rhs has type %a" - (Pretty_printing.pretty_print_assignmentoperator assignop) - UnsizedType.pp lt UnsizedType.pp rt + Fmt.pf + ppf + "Ill-typed arguments supplied to assignment operator %s: lhs has type %a and rhs \ + has type %a" + (Pretty_printing.pretty_print_assignmentoperator assignop) + UnsizedType.pp + lt + UnsizedType.pp + rt | IllTypedTernaryIf (UInt, ut2, ut3) -> - Fmt.pf ppf - "Type mismatch in ternary expression, expression when true is: %a; \ - expression when false is: %a" - UnsizedType.pp ut2 UnsizedType.pp ut3 + Fmt.pf + ppf + "Type mismatch in ternary expression, expression when true is: %a; expression \ + when false is: %a" + UnsizedType.pp + ut2 + UnsizedType.pp + ut3 | IllTypedTernaryIf (ut1, _, _) -> - Fmt.pf ppf - "Condition in ternary expression must be primitive int; found type=%a" - UnsizedType.pp ut1 + Fmt.pf + ppf + "Condition in ternary expression must be primitive int; found type=%a" + UnsizedType.pp + ut1 | IllTypedReduceSum (name, arg_tys, args) -> - let arg_types = List.map ~f:(fun (_, t) -> t) args in - let first, rest = List.split_n arg_types 1 in - let generate_reduce_sum_sig = - List.concat - [ [ UnsizedType.UFun - ( List.hd_exn args :: (AutoDiffable, UInt) - :: (AutoDiffable, UInt) :: List.tl_exn args - , ReturnType UReal ) ] - ; first; [UInt]; rest ] - in - Fmt.pf ppf - "Ill-typed arguments supplied to function '%s'. Expected \ - arguments:@[%a@]\n\ - @[Instead supplied arguments of incompatible type: %a@]" - name - Fmt.(list UnsizedType.pp ~sep:comma) - generate_reduce_sum_sig - Fmt.(list UnsizedType.pp ~sep:comma) - arg_tys + let arg_types = List.map ~f:(fun (_, t) -> t) args in + let first, rest = List.split_n arg_types 1 in + let generate_reduce_sum_sig = + List.concat + [ [ UnsizedType.UFun + ( List.hd_exn args + :: (AutoDiffable, UInt) + :: (AutoDiffable, UInt) + :: List.tl_exn args + , ReturnType UReal ) + ] + ; first + ; [ UInt ] + ; rest + ] + in + Fmt.pf + ppf + "Ill-typed arguments supplied to function '%s'. Expected arguments:@[%a@]\n\ + @[Instead supplied arguments of incompatible type: %a@]" + name + Fmt.(list UnsizedType.pp ~sep:comma) + generate_reduce_sum_sig + Fmt.(list UnsizedType.pp ~sep:comma) + arg_tys | IllTypedReduceSumGeneric (name, arg_tys) -> - let rec n_commas n = if n = 0 then "" else "," ^ n_commas (n - 1) in - let type_string (a, b, c, d) i = - Fmt.strf "(T[%s], %a, %a, ...) => %a, T[%s], %a, ...\n" - (n_commas (i - 1)) - Pretty_printing.pp_unsizedtype a Pretty_printing.pp_unsizedtype b - Pretty_printing.pp_unsizedtype c - (n_commas (i - 1)) - Pretty_printing.pp_unsizedtype d - in - let lines = - List.map - ~f:(fun i -> type_string (UInt, UInt, UReal, UInt) i) - Stan_math_signatures.reduce_sum_allowed_dimensionalities - in - Fmt.pf ppf - "Ill-typed arguments supplied to function '%s'. Available arguments:\n\ - %sWhere T is any one of int, real, vector, row_vector or \ - matrix.@[Instead supplied arguments of incompatible type: %a@]" - name - (String.concat ~sep:"" lines) - Fmt.(list UnsizedType.pp ~sep:comma) - arg_tys + let rec n_commas n = if n = 0 then "" else "," ^ n_commas (n - 1) in + let type_string (a, b, c, d) i = + Fmt.strf + "(T[%s], %a, %a, ...) => %a, T[%s], %a, ...\n" + (n_commas (i - 1)) + Pretty_printing.pp_unsizedtype + a + Pretty_printing.pp_unsizedtype + b + Pretty_printing.pp_unsizedtype + c + (n_commas (i - 1)) + Pretty_printing.pp_unsizedtype + d + in + let lines = + List.map + ~f:(fun i -> type_string (UInt, UInt, UReal, UInt) i) + Stan_math_signatures.reduce_sum_allowed_dimensionalities + in + Fmt.pf + ppf + "Ill-typed arguments supplied to function '%s'. Available arguments:\n\ + %sWhere T is any one of int, real, vector, row_vector or matrix.@[Instead \ + supplied arguments of incompatible type: %a@]" + name + (String.concat ~sep:"" lines) + Fmt.(list UnsizedType.pp ~sep:comma) + arg_tys | NotIndexable ut -> - Fmt.pf ppf - "Only expressions of array, matrix, row_vector and vector type may \ - be indexed. Instead, found type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "Only expressions of array, matrix, row_vector and vector type may be indexed. \ + Instead, found type %a." + UnsizedType.pp + ut | ReturningFnExpectedNonReturningFound fn_name -> - Fmt.pf ppf - "A returning function was expected but a non-returning function \ - '%s' was supplied." - fn_name + Fmt.pf + ppf + "A returning function was expected but a non-returning function '%s' was \ + supplied." + fn_name | NonReturningFnExpectedReturningFound fn_name -> - Fmt.pf ppf - "A non-returning function was expected but a returning function \ - '%s' was supplied." - fn_name + Fmt.pf + ppf + "A non-returning function was expected but a returning function '%s' was \ + supplied." + fn_name | ReturningFnExpectedNonFnFound fn_name -> - Fmt.pf ppf - "A returning function was expected but a non-function value '%s' \ - was supplied." - fn_name + Fmt.pf + ppf + "A returning function was expected but a non-function value '%s' was supplied." + fn_name | NonReturningFnExpectedNonFnFound fn_name -> - Fmt.pf ppf - "A non-returning function was expected but a non-function value \ - '%s' was supplied." - fn_name + Fmt.pf + ppf + "A non-returning function was expected but a non-function value '%s' was \ + supplied." + fn_name | ReturningFnExpectedUndeclaredIdentFound fn_name -> - Fmt.pf ppf - "A returning function was expected but an undeclared identifier \ - '%s' was supplied." - fn_name + Fmt.pf + ppf + "A returning function was expected but an undeclared identifier '%s' was \ + supplied." + fn_name | NonReturningFnExpectedUndeclaredIdentFound fn_name -> - Fmt.pf ppf - "A non-returning function was expected but an undeclared identifier \ - '%s' was supplied." - fn_name + Fmt.pf + ppf + "A non-returning function was expected but an undeclared identifier '%s' was \ + supplied." + fn_name | IllTypedStanLibFunctionApp (name, arg_tys) -> - Fmt.pf ppf - "Ill-typed arguments supplied to function '%s'. Available \ - signatures: %s@[Instead supplied arguments of incompatible \ - type: %a.@]" - name - (Stan_math_signatures.pretty_print_math_sigs name) - Fmt.(list UnsizedType.pp ~sep:comma) - arg_tys + Fmt.pf + ppf + "Ill-typed arguments supplied to function '%s'. Available signatures: \ + %s@[Instead supplied arguments of incompatible type: %a.@]" + name + (Stan_math_signatures.pretty_print_math_sigs name) + Fmt.(list UnsizedType.pp ~sep:comma) + arg_tys | IllTypedUserDefinedFunctionApp (name, listed_tys, return_ty, arg_tys) -> - Fmt.pf ppf - "Ill-typed arguments supplied to function '%s'. Available \ - signatures:%a\n\ - @[Instead supplied arguments of incompatible type: %a.@]" - name UnsizedType.pp - (UFun (listed_tys, return_ty)) - Fmt.(list UnsizedType.pp ~sep:comma) - arg_tys + Fmt.pf + ppf + "Ill-typed arguments supplied to function '%s'. Available signatures:%a\n\ + @[Instead supplied arguments of incompatible type: %a.@]" + name + UnsizedType.pp + (UFun (listed_tys, return_ty)) + Fmt.(list UnsizedType.pp ~sep:comma) + arg_tys | IllTypedBinaryOperator (op, lt, rt) -> - Fmt.pf ppf - "Ill-typed arguments supplied to infix operator %a. Available \ - signatures: %s@[Instead supplied arguments of incompatible \ - type: %a, %a.@]" - Operator.pp op - ( Stan_math_signatures.pretty_print_math_lib_operator_sigs op - |> String.concat ~sep:"\n" ) - UnsizedType.pp lt UnsizedType.pp rt + Fmt.pf + ppf + "Ill-typed arguments supplied to infix operator %a. Available signatures: \ + %s@[Instead supplied arguments of incompatible type: %a, %a.@]" + Operator.pp + op + (Stan_math_signatures.pretty_print_math_lib_operator_sigs op + |> String.concat ~sep:"\n") + UnsizedType.pp + lt + UnsizedType.pp + rt | IllTypedPrefixOperator (op, ut) -> - Fmt.pf ppf - "Ill-typed arguments supplied to prefix operator %a. Available \ - signatures: %s@[Instead supplied argument of incompatible type: \ - %a.@]" - Operator.pp op - ( Stan_math_signatures.pretty_print_math_lib_operator_sigs op - |> String.concat ~sep:"\n" ) - UnsizedType.pp ut + Fmt.pf + ppf + "Ill-typed arguments supplied to prefix operator %a. Available signatures: \ + %s@[Instead supplied argument of incompatible type: %a.@]" + Operator.pp + op + (Stan_math_signatures.pretty_print_math_lib_operator_sigs op + |> String.concat ~sep:"\n") + UnsizedType.pp + ut | IllTypedPostfixOperator (op, ut) -> - Fmt.pf ppf - "Ill-typed arguments supplied to postfix operator %a. Available \ - signatures: %s\n\ - Instead supplied argument of incompatible type: %a." - Operator.pp op - ( Stan_math_signatures.pretty_print_math_lib_operator_sigs op - |> String.concat ~sep:"\n" ) - UnsizedType.pp ut + Fmt.pf + ppf + "Ill-typed arguments supplied to postfix operator %a. Available signatures: %s\n\ + Instead supplied argument of incompatible type: %a." + Operator.pp + op + (Stan_math_signatures.pretty_print_math_lib_operator_sigs op + |> String.concat ~sep:"\n") + UnsizedType.pp + ut + ;; end module IdentifierError = struct @@ -244,14 +307,12 @@ module IdentifierError = struct let pp ppf = function | IsStanMathName name -> - Fmt.pf ppf "Identifier '%s' clashes with Stan Math library function." - name + Fmt.pf ppf "Identifier '%s' clashes with Stan Math library function." name | InUse name -> Fmt.pf ppf "Identifier '%s' is already in use." name - | IsModelName name -> - Fmt.pf ppf "Identifier '%s' clashes with model name." name - | IsKeyword name -> - Fmt.pf ppf "Identifier '%s' clashes with reserved keyword." name + | IsModelName name -> Fmt.pf ppf "Identifier '%s' clashes with model name." name + | IsKeyword name -> Fmt.pf ppf "Identifier '%s' clashes with reserved keyword." name | NotInScope name -> Fmt.pf ppf "Identifier '%s' not in scope." name + ;; end module ExpressionError = struct @@ -267,32 +328,33 @@ module ExpressionError = struct let pp ppf = function | InvalidMapRectFn fn_name -> - Fmt.pf ppf - "Mapped function cannot be an _rng or _lp function, found function \ - name: %s" - fn_name + Fmt.pf + ppf + "Mapped function cannot be an _rng or _lp function, found function name: %s" + fn_name | InvalidSizeDeclRng -> - Fmt.pf ppf - "Random number generators are not allowed in top level size \ - declarations." + Fmt.pf + ppf + "Random number generators are not allowed in top level size declarations." | InvalidRngFunction -> - Fmt.pf ppf - "Random number generators are only allowed in transformed data \ - block, generated quantities block or user-defined functions with \ - names ending in _rng." + Fmt.pf + ppf + "Random number generators are only allowed in transformed data block, generated \ + quantities block or user-defined functions with names ending in _rng." | ConditionalNotationNotAllowed -> - Fmt.pf ppf - "Only functions with names ending in _lpdf, _lpmf, _lcdf, _lccdf \ - can make use of conditional notation." + Fmt.pf + ppf + "Only functions with names ending in _lpdf, _lpmf, _lcdf, _lccdf can make use of \ + conditional notation." | ConditioningRequired -> - Fmt.pf ppf - "Probabilty functions with suffixes _lpdf, _lpmf, _lcdf, and \ - _lccdf, require a vertical bar (|) between the first two arguments." + Fmt.pf + ppf + "Probabilty functions with suffixes _lpdf, _lpmf, _lcdf, and _lccdf, require a \ + vertical bar (|) between the first two arguments." | NotPrintable -> Fmt.pf ppf "Functions cannot be printed." - | EmptyArray -> - Fmt.pf ppf "Array expressions must contain at least one element." - | IntTooLarge -> - Fmt.pf ppf "Integer literal cannot be larger than 2_147_483_647." + | EmptyArray -> Fmt.pf ppf "Array expressions must contain at least one element." + | IntTooLarge -> Fmt.pf ppf "Integer literal cannot be larger than 2_147_483_647." + ;; end module StatementError = struct @@ -324,102 +386,105 @@ module StatementError = struct let pp ppf = function | CannotAssignToReadOnly name -> - Fmt.pf ppf - "Cannot assign to function argument or loop identifier '%s'." name + Fmt.pf ppf "Cannot assign to function argument or loop identifier '%s'." name | CannotAssignToGlobal name -> - Fmt.pf ppf - "Cannot assign to global variable '%s' declared in previous blocks." - name + Fmt.pf ppf "Cannot assign to global variable '%s' declared in previous blocks." name | TargetPlusEqualsOutsideModelOrLogProb -> - Fmt.pf ppf - "Target can only be accessed in the model block or in definitions \ - of functions with the suffix _lp." + Fmt.pf + ppf + "Target can only be accessed in the model block or in definitions of functions \ + with the suffix _lp." | InvalidSamplingPDForPMF -> - Fmt.pf ppf - {| + Fmt.pf + ppf + {| ~ statement should refer to a distribution without its "_lpdf" or "_lpmf" suffix. For example, "target += normal_lpdf(y, 0, 1)" should become "y ~ normal(0, 1)." |} | InvalidSamplingCDForCCDF name -> - Fmt.pf ppf - "CDF and CCDF functions may not be used with sampling notation. Use \ - increment_log_prob(%s_log(...)) instead." - name + Fmt.pf + ppf + "CDF and CCDF functions may not be used with sampling notation. Use \ + increment_log_prob(%s_log(...)) instead." + name | InvalidSamplingNoSuchDistribution name -> - Fmt.pf ppf - "Ill-typed arguments to '~' statement. No distribution '%s' was \ - found with the correct signature." - name + Fmt.pf + ppf + "Ill-typed arguments to '~' statement. No distribution '%s' was found with the \ + correct signature." + name | InvalidTruncationCDForCCDF -> - Fmt.pf ppf - "Truncation is only defined if distribution has _lcdf and _lccdf \ - functions implemented with appropriate signature." + Fmt.pf + ppf + "Truncation is only defined if distribution has _lcdf and _lccdf functions \ + implemented with appropriate signature." | MultivariateTruncation -> - Fmt.pf ppf "Outcomes in truncated distributions must be univariate." - | BreakOutsideLoop -> - Fmt.pf ppf "Break statements may only be used in loops." - | ContinueOutsideLoop -> - Fmt.pf ppf "Continue statements may only be used in loops." + Fmt.pf ppf "Outcomes in truncated distributions must be univariate." + | BreakOutsideLoop -> Fmt.pf ppf "Break statements may only be used in loops." + | ContinueOutsideLoop -> Fmt.pf ppf "Continue statements may only be used in loops." | ExpressionReturnOutsideReturningFn -> - Fmt.pf ppf - "Expression return statements may only be used inside returning \ - function definitions." + Fmt.pf + ppf + "Expression return statements may only be used inside returning function \ + definitions." | VoidReturnOutsideNonReturningFn -> - Fmt.pf ppf - "Void return statements may only be used inside non-returning \ - function definitions." + Fmt.pf + ppf + "Void return statements may only be used inside non-returning function \ + definitions." | NonDataVariableSizeDecl -> - Fmt.pf ppf - "Non-data variables are not allowed in top level size declarations." + Fmt.pf ppf "Non-data variables are not allowed in top level size declarations." | NonIntBounds -> - Fmt.pf ppf - "Bounds of integer variable must be of type int. Found type real." - | TransformedParamsInt -> - Fmt.pf ppf "(Transformed) Parameters cannot be integers." + Fmt.pf ppf "Bounds of integer variable must be of type int. Found type real." + | TransformedParamsInt -> Fmt.pf ppf "(Transformed) Parameters cannot be integers." | MismatchFunDefDecl (name, Some ut) -> - Fmt.pf ppf "Function '%s' has already been declared to have type %a" - name UnsizedType.pp ut + Fmt.pf + ppf + "Function '%s' has already been declared to have type %a" + name + UnsizedType.pp + ut | MismatchFunDefDecl (name, None) -> - Fmt.pf ppf - "Function '%s' has already been declared but type cannot be \ - determined." - name + Fmt.pf + ppf + "Function '%s' has already been declared but type cannot be determined." + name | FunDeclExists name -> - Fmt.pf ppf - "Function '%s' has already been declared. A definition is expected." - name + Fmt.pf ppf "Function '%s' has already been declared. A definition is expected." name | FunDeclNoDefn -> - Fmt.pf ppf "Some function is declared without specifying a definition." + Fmt.pf ppf "Some function is declared without specifying a definition." | FunDeclNeedsBlock -> - Fmt.pf ppf "Function definitions must be wrapped in curly braces." + Fmt.pf ppf "Function definitions must be wrapped in curly braces." | NonRealProbFunDef -> - Fmt.pf ppf - "Real return type required for probability functions ending in \ - _log, _lpdf, _lpmf, _lcdf, or _lccdf." + Fmt.pf + ppf + "Real return type required for probability functions ending in _log, _lpdf, \ + _lpmf, _lcdf, or _lccdf." | ProbDensityNonRealVariate (Some ut) -> - Fmt.pf ppf - "Probability density functions require real variates (first \ - argument). Instead found type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "Probability density functions require real variates (first argument). Instead \ + found type %a." + UnsizedType.pp + ut | ProbDensityNonRealVariate _ -> - Fmt.pf ppf - "Probability density functions require real variates (first \ - argument)." + Fmt.pf ppf "Probability density functions require real variates (first argument)." | ProbMassNonIntVariate (Some ut) -> - Fmt.pf ppf - "Probability mass functions require integer variates (first \ - argument). Instead found type %a." - UnsizedType.pp ut + Fmt.pf + ppf + "Probability mass functions require integer variates (first argument). Instead \ + found type %a." + UnsizedType.pp + ut | ProbMassNonIntVariate _ -> - Fmt.pf ppf - "Probability mass functions require integer variates (first \ - argument)." + Fmt.pf ppf "Probability mass functions require integer variates (first argument)." | DuplicateArgNames -> - Fmt.pf ppf "All function arguments must have distinct identifiers." + Fmt.pf ppf "All function arguments must have distinct identifiers." | IncompatibleReturnType -> - Fmt.pf ppf - "Function bodies must contain a return statement of correct type in \ - every branch." + Fmt.pf + ppf + "Function bodies must contain a return statement of correct type in every branch." + ;; end type t = @@ -433,193 +498,208 @@ let pp ppf = function | IdentifierError (_, err) -> IdentifierError.pp ppf err | ExpressionError (_, err) -> ExpressionError.pp ppf err | StatementError (_, err) -> StatementError.pp ppf err +;; let location = function | TypeError (loc, _) -> loc | IdentifierError (loc, _) -> loc | ExpressionError (loc, _) -> loc | StatementError (loc, _) -> loc +;; (* -- Constructors ---------------------------------------------------------- *) let mismatched_return_types loc rt1 rt2 = TypeError (loc, TypeError.MismatchedReturnTypes (rt1, rt2)) +;; let mismatched_array_types loc t1 t2 = TypeError (loc, TypeError.MismatchedArrayTypes (t1, t2)) +;; -let invalid_row_vector_types loc ty = - TypeError (loc, TypeError.InvalidRowVectorTypes ty) - -let invalid_matrix_types loc ty = - TypeError (loc, TypeError.InvalidMatrixTypes ty) - +let invalid_row_vector_types loc ty = TypeError (loc, TypeError.InvalidRowVectorTypes ty) +let invalid_matrix_types loc ty = TypeError (loc, TypeError.InvalidMatrixTypes ty) let int_expected loc name ut = TypeError (loc, TypeError.IntExpected (name, ut)) let int_or_real_expected loc name ut = TypeError (loc, TypeError.IntOrRealExpected (name, ut)) +;; let scalar_or_type_expected loc name et ut = TypeError (loc, TypeError.TypeExpected (name, et, ut)) +;; let int_intarray_or_range_expected loc ut = TypeError (loc, TypeError.IntIntArrayOrRangeExpected ut) +;; let int_or_real_container_expected loc ut = TypeError (loc, TypeError.IntOrRealContainerExpected ut) +;; let array_vector_rowvector_matrix_expected loc ut = TypeError (loc, TypeError.ArrayVectorRowVectorMatrixExpected ut) +;; let illtyped_assignment loc assignop lt rt = TypeError (loc, TypeError.IllTypedAssignment (assignop, lt, rt)) +;; let illtyped_ternary_if loc predt lt rt = TypeError (loc, TypeError.IllTypedTernaryIf (predt, lt, rt)) +;; let returning_fn_expected_nonreturning_found loc name = TypeError (loc, TypeError.ReturningFnExpectedNonReturningFound name) +;; let illtyped_reduce_sum loc name arg_tys args = TypeError (loc, TypeError.IllTypedReduceSum (name, arg_tys, args)) +;; let illtyped_reduce_sum_generic loc name arg_tys = TypeError (loc, TypeError.IllTypedReduceSumGeneric (name, arg_tys)) +;; let returning_fn_expected_nonfn_found loc name = TypeError (loc, TypeError.ReturningFnExpectedNonFnFound name) +;; let returning_fn_expected_undeclaredident_found loc name = TypeError (loc, TypeError.ReturningFnExpectedUndeclaredIdentFound name) +;; let nonreturning_fn_expected_returning_found loc name = TypeError (loc, TypeError.NonReturningFnExpectedReturningFound name) +;; let nonreturning_fn_expected_nonfn_found loc name = TypeError (loc, TypeError.NonReturningFnExpectedNonFnFound name) +;; let nonreturning_fn_expected_undeclaredident_found loc name = TypeError (loc, TypeError.NonReturningFnExpectedUndeclaredIdentFound name) +;; let illtyped_stanlib_fn_app loc name arg_tys = TypeError (loc, TypeError.IllTypedStanLibFunctionApp (name, arg_tys)) +;; let illtyped_userdefined_fn_app loc name decl_arg_tys decl_return_ty arg_tys = TypeError ( loc , TypeError.IllTypedUserDefinedFunctionApp (name, decl_arg_tys, decl_return_ty, arg_tys) ) +;; let illtyped_binary_op loc op lt rt = TypeError (loc, TypeError.IllTypedBinaryOperator (op, lt, rt)) +;; let illtyped_prefix_op loc op ut = TypeError (loc, TypeError.IllTypedPrefixOperator (op, ut)) +;; let illtyped_postfix_op loc op ut = TypeError (loc, TypeError.IllTypedPostfixOperator (op, ut)) +;; let not_indexable loc ut = TypeError (loc, TypeError.NotIndexable ut) - -let ident_is_keyword loc name = - IdentifierError (loc, IdentifierError.IsKeyword name) - -let ident_is_model_name loc name = - IdentifierError (loc, IdentifierError.IsModelName name) +let ident_is_keyword loc name = IdentifierError (loc, IdentifierError.IsKeyword name) +let ident_is_model_name loc name = IdentifierError (loc, IdentifierError.IsModelName name) let ident_is_stanmath_name loc name = IdentifierError (loc, IdentifierError.IsStanMathName name) +;; let ident_in_use loc name = IdentifierError (loc, IdentifierError.InUse name) - -let ident_not_in_scope loc name = - IdentifierError (loc, IdentifierError.NotInScope name) +let ident_not_in_scope loc name = IdentifierError (loc, IdentifierError.NotInScope name) let invalid_map_rect_fn loc name = ExpressionError (loc, ExpressionError.InvalidMapRectFn name) +;; -let invalid_decl_rng_fn loc = - ExpressionError (loc, ExpressionError.InvalidSizeDeclRng) - -let invalid_rng_fn loc = - ExpressionError (loc, ExpressionError.InvalidRngFunction) +let invalid_decl_rng_fn loc = ExpressionError (loc, ExpressionError.InvalidSizeDeclRng) +let invalid_rng_fn loc = ExpressionError (loc, ExpressionError.InvalidRngFunction) let conditional_notation_not_allowed loc = ExpressionError (loc, ExpressionError.ConditionalNotationNotAllowed) +;; -let conditioning_required loc = - ExpressionError (loc, ExpressionError.ConditioningRequired) - +let conditioning_required loc = ExpressionError (loc, ExpressionError.ConditioningRequired) let not_printable loc = ExpressionError (loc, ExpressionError.NotPrintable) let empty_array loc = ExpressionError (loc, ExpressionError.EmptyArray) let bad_int_literal loc = ExpressionError (loc, ExpressionError.IntTooLarge) let cannot_assign_to_read_only loc name = StatementError (loc, StatementError.CannotAssignToReadOnly name) +;; let cannot_assign_to_global loc name = StatementError (loc, StatementError.CannotAssignToGlobal name) +;; let invalid_sampling_pdf_or_pmf loc = StatementError (loc, StatementError.InvalidSamplingPDForPMF) +;; let invalid_sampling_cdf_or_ccdf loc name = StatementError (loc, StatementError.InvalidSamplingCDForCCDF name) +;; let invalid_sampling_no_such_dist loc name = StatementError (loc, StatementError.InvalidSamplingNoSuchDistribution name) +;; let target_plusequals_outisde_model_or_logprob loc = StatementError (loc, StatementError.TargetPlusEqualsOutsideModelOrLogProb) +;; let invalid_truncation_cdf_or_ccdf loc = StatementError (loc, StatementError.InvalidTruncationCDForCCDF) +;; let multivariate_truncation loc = StatementError (loc, StatementError.MultivariateTruncation) +;; -let break_outside_loop loc = - StatementError (loc, StatementError.BreakOutsideLoop) - -let continue_outside_loop loc = - StatementError (loc, StatementError.ContinueOutsideLoop) +let break_outside_loop loc = StatementError (loc, StatementError.BreakOutsideLoop) +let continue_outside_loop loc = StatementError (loc, StatementError.ContinueOutsideLoop) let expression_return_outside_returning_fn loc = StatementError (loc, StatementError.ExpressionReturnOutsideReturningFn) +;; let void_ouside_nonreturning_fn loc = StatementError (loc, StatementError.VoidReturnOutsideNonReturningFn) +;; let non_data_variable_size_decl loc = StatementError (loc, StatementError.NonDataVariableSizeDecl) +;; let non_int_bounds loc = StatementError (loc, StatementError.NonIntBounds) - -let transformed_params_int loc = - StatementError (loc, StatementError.TransformedParamsInt) +let transformed_params_int loc = StatementError (loc, StatementError.TransformedParamsInt) let mismatched_fn_def_decl loc name ut_opt = StatementError (loc, StatementError.MismatchFunDefDecl (name, ut_opt)) +;; -let fn_decl_exists loc name = - StatementError (loc, StatementError.FunDeclExists name) - +let fn_decl_exists loc name = StatementError (loc, StatementError.FunDeclExists name) let fn_decl_without_def loc = StatementError (loc, StatementError.FunDeclNoDefn) - -let fn_decl_needs_block loc = - StatementError (loc, StatementError.FunDeclNeedsBlock) - -let non_real_prob_fn_def loc = - StatementError (loc, StatementError.NonRealProbFunDef) +let fn_decl_needs_block loc = StatementError (loc, StatementError.FunDeclNeedsBlock) +let non_real_prob_fn_def loc = StatementError (loc, StatementError.NonRealProbFunDef) let prob_density_non_real_variate loc ut_opt = StatementError (loc, StatementError.ProbDensityNonRealVariate ut_opt) +;; let prob_mass_non_int_variate loc ut_opt = StatementError (loc, StatementError.ProbMassNonIntVariate ut_opt) +;; -let duplicate_arg_names loc = - StatementError (loc, StatementError.DuplicateArgNames) +let duplicate_arg_names loc = StatementError (loc, StatementError.DuplicateArgNames) let incompatible_return_types loc = StatementError (loc, StatementError.IncompatibleReturnType) +;; diff --git a/src/frontend/Semantic_error.mli b/src/frontend/Semantic_error.mli index 598365e0b3..1e8c657497 100644 --- a/src/frontend/Semantic_error.mli +++ b/src/frontend/Semantic_error.mli @@ -5,12 +5,13 @@ type t val pp : Format.formatter -> t -> unit val location : t -> Location_span.t -val mismatched_return_types : - Location_span.t -> UnsizedType.returntype -> UnsizedType.returntype -> t - -val mismatched_array_types : - Location_span.t -> UnsizedType.t -> UnsizedType.t -> t +val mismatched_return_types + : Location_span.t + -> UnsizedType.returntype + -> UnsizedType.returntype + -> t +val mismatched_array_types : Location_span.t -> UnsizedType.t -> UnsizedType.t -> t val invalid_row_vector_types : Location_span.t -> UnsizedType.t -> t val invalid_matrix_types : Location_span.t -> UnsizedType.t -> t val int_expected : Location_span.t -> string -> UnsizedType.t -> t @@ -18,57 +19,60 @@ val int_or_real_expected : Location_span.t -> string -> UnsizedType.t -> t val int_intarray_or_range_expected : Location_span.t -> UnsizedType.t -> t val int_or_real_container_expected : Location_span.t -> UnsizedType.t -> t -val scalar_or_type_expected : - Location_span.t -> string -> UnsizedType.t -> UnsizedType.t -> t +val scalar_or_type_expected + : Location_span.t + -> string + -> UnsizedType.t + -> UnsizedType.t + -> t -val array_vector_rowvector_matrix_expected : - Location_span.t -> UnsizedType.t -> t +val array_vector_rowvector_matrix_expected : Location_span.t -> UnsizedType.t -> t -val illtyped_assignment : - Location_span.t +val illtyped_assignment + : Location_span.t -> Ast.assignmentoperator -> UnsizedType.t -> UnsizedType.t -> t -val illtyped_ternary_if : - Location_span.t -> UnsizedType.t -> UnsizedType.t -> UnsizedType.t -> t +val illtyped_ternary_if + : Location_span.t + -> UnsizedType.t + -> UnsizedType.t + -> UnsizedType.t + -> t val returning_fn_expected_nonreturning_found : Location_span.t -> string -> t val returning_fn_expected_nonfn_found : Location_span.t -> string -> t +val returning_fn_expected_undeclaredident_found : Location_span.t -> string -> t -val returning_fn_expected_undeclaredident_found : - Location_span.t -> string -> t - -val illtyped_reduce_sum : - Location_span.t +val illtyped_reduce_sum + : Location_span.t -> string -> UnsizedType.t list -> (UnsizedType.autodifftype * UnsizedType.t) list -> t -val illtyped_reduce_sum_generic : - Location_span.t -> string -> UnsizedType.t list -> t - +val illtyped_reduce_sum_generic : Location_span.t -> string -> UnsizedType.t list -> t val nonreturning_fn_expected_returning_found : Location_span.t -> string -> t val nonreturning_fn_expected_nonfn_found : Location_span.t -> string -> t +val nonreturning_fn_expected_undeclaredident_found : Location_span.t -> string -> t +val illtyped_stanlib_fn_app : Location_span.t -> string -> UnsizedType.t list -> t -val nonreturning_fn_expected_undeclaredident_found : - Location_span.t -> string -> t - -val illtyped_stanlib_fn_app : - Location_span.t -> string -> UnsizedType.t list -> t - -val illtyped_userdefined_fn_app : - Location_span.t +val illtyped_userdefined_fn_app + : Location_span.t -> string -> (UnsizedType.autodifftype * UnsizedType.t) list -> UnsizedType.returntype -> UnsizedType.t list -> t -val illtyped_binary_op : - Location_span.t -> Operator.t -> UnsizedType.t -> UnsizedType.t -> t +val illtyped_binary_op + : Location_span.t + -> Operator.t + -> UnsizedType.t + -> UnsizedType.t + -> t val illtyped_prefix_op : Location_span.t -> Operator.t -> UnsizedType.t -> t val illtyped_postfix_op : Location_span.t -> Operator.t -> UnsizedType.t -> t @@ -101,18 +105,12 @@ val void_ouside_nonreturning_fn : Location_span.t -> t val non_data_variable_size_decl : Location_span.t -> t val non_int_bounds : Location_span.t -> t val transformed_params_int : Location_span.t -> t - -val mismatched_fn_def_decl : - Location_span.t -> string -> UnsizedType.t option -> t - +val mismatched_fn_def_decl : Location_span.t -> string -> UnsizedType.t option -> t val fn_decl_exists : Location_span.t -> string -> t val fn_decl_without_def : Location_span.t -> t val fn_decl_needs_block : Location_span.t -> t val non_real_prob_fn_def : Location_span.t -> t - -val prob_density_non_real_variate : - Location_span.t -> UnsizedType.t option -> t - +val prob_density_non_real_variate : Location_span.t -> UnsizedType.t option -> t val prob_mass_non_int_variate : Location_span.t -> UnsizedType.t option -> t val duplicate_arg_names : Location_span.t -> t val incompatible_return_types : Location_span.t -> t diff --git a/src/frontend/Symbol_table.ml b/src/frontend/Symbol_table.ml index a6d975ae21..cfc6b05a24 100644 --- a/src/frontend/Symbol_table.ml +++ b/src/frontend/Symbol_table.ml @@ -5,75 +5,90 @@ open Core_kernel (* TODO: I'm sure this implementation could be made more efficient if that's necessary. There's no need for all the string comparison. We could just keep track of the count of the entry into the hash table and use that for comparison. *) type 'a state = - { table: (string, 'a) Hashtbl.t - ; stack: string Stack.t - ; scopedepth: int ref - ; readonly: (string, unit) Hashtbl.t - ; isunassigned: (string, unit) Hashtbl.t - ; globals: (string, unit) Hashtbl.t } + { table : (string, 'a) Hashtbl.t + ; stack : string Stack.t + ; scopedepth : int ref + ; readonly : (string, unit) Hashtbl.t + ; isunassigned : (string, unit) Hashtbl.t + ; globals : (string, unit) Hashtbl.t + } let initialize () = - { table= String.Table.create () - ; stack= Stack.create () - ; scopedepth= ref 0 - ; readonly= String.Table.create () - ; isunassigned= String.Table.create () - ; globals= String.Table.create () } + { table = String.Table.create () + ; stack = Stack.create () + ; scopedepth = ref 0 + ; readonly = String.Table.create () + ; isunassigned = String.Table.create () + ; globals = String.Table.create () + } +;; let enter s str ty = - let _ : [`Duplicate | `Ok] = - if !(s.scopedepth) = 0 then Hashtbl.add s.globals ~key:str ~data:() - else `Ok + let (_ : [ `Duplicate | `Ok ]) = + if !(s.scopedepth) = 0 then Hashtbl.add s.globals ~key:str ~data:() else `Ok in - let _ : [`Duplicate | `Ok] = Hashtbl.add s.table ~key:str ~data:ty in + let (_ : [ `Duplicate | `Ok ]) = Hashtbl.add s.table ~key:str ~data:ty in Stack.push s.stack str +;; let look s str = Hashtbl.find s.table str let begin_scope s = - s.scopedepth := !(s.scopedepth) + 1 ; + s.scopedepth := !(s.scopedepth) + 1; Stack.push s.stack "-sentinel-new-scope-" +;; (* using a string "-sentinel-new-scope-" here that can never be used as an identifier to indicate that new scope is entered *) let end_scope s = - s.scopedepth := !(s.scopedepth) - 1 ; + s.scopedepth := !(s.scopedepth) - 1; while Stack.top_exn s.stack <> "-sentinel-new-scope-" do (* we pop the stack down to where we entered the current scope and remove all variables defined since from the var map *) - Hashtbl.remove s.table (Stack.top_exn s.stack) ; - Hashtbl.remove s.readonly (Stack.top_exn s.stack) ; - Hashtbl.remove s.isunassigned (Stack.top_exn s.stack) ; - let _ : string = Stack.pop_exn s.stack in + Hashtbl.remove s.table (Stack.top_exn s.stack); + Hashtbl.remove s.readonly (Stack.top_exn s.stack); + Hashtbl.remove s.isunassigned (Stack.top_exn s.stack); + let (_ : string) = Stack.pop_exn s.stack in () - done ; - let _ : string = Stack.pop_exn s.stack in + done; + let (_ : string) = Stack.pop_exn s.stack in () +;; let set_read_only s str = - let _ : [`Duplicate | `Ok] = Hashtbl.add s.readonly ~key:str ~data:() in + let (_ : [ `Duplicate | `Ok ]) = Hashtbl.add s.readonly ~key:str ~data:() in () +;; let get_read_only s str = - match Hashtbl.find s.readonly str with Some () -> true | _ -> false + match Hashtbl.find s.readonly str with + | Some () -> true + | _ -> false +;; let set_is_assigned s str = Hashtbl.remove s.isunassigned str let set_is_unassigned s str = - let _ : [`Duplicate | `Ok] = - if Hashtbl.mem s.isunassigned str then `Ok + let (_ : [ `Duplicate | `Ok ]) = + if Hashtbl.mem s.isunassigned str + then `Ok else Hashtbl.add s.isunassigned ~key:str ~data:() in () +;; let check_is_unassigned s str = Hashtbl.mem s.isunassigned str let check_some_id_is_unassigned s = not (Hashtbl.length s.isunassigned = 0) let is_global s str = - match Hashtbl.find s.globals str with Some _ -> true | _ -> false + match Hashtbl.find s.globals str with + | Some _ -> true + | _ -> false +;; let unsafe_clear_symbol_table s = - Hashtbl.clear s.table ; - Stack.clear s.stack ; - s.scopedepth := 0 ; - Hashtbl.clear s.readonly ; - Hashtbl.clear s.isunassigned ; + Hashtbl.clear s.table; + Stack.clear s.stack; + s.scopedepth := 0; + Hashtbl.clear s.readonly; + Hashtbl.clear s.isunassigned; Hashtbl.clear s.globals +;; diff --git a/src/frontend/Symbol_table.mli b/src/frontend/Symbol_table.mli index 60dae7583a..ae7c5731ec 100644 --- a/src/frontend/Symbol_table.mli +++ b/src/frontend/Symbol_table.mli @@ -2,42 +2,42 @@ type 'a state -val initialize : unit -> 'a state (** Creates a new symbol table *) +val initialize : unit -> 'a state -val enter : 'a state -> string -> 'a -> unit (** Enters a specified identifier with its specified type (or other) information into a symbol table *) +val enter : 'a state -> string -> 'a -> unit -val look : 'a state -> string -> 'a option (** Looks for an identifier in a symbol table and returns its information if found and None otherwise *) +val look : 'a state -> string -> 'a option -val begin_scope : 'a state -> unit (** Used to start a new local scope which symbols added from now will end up in *) +val begin_scope : 'a state -> unit -val end_scope : 'a state -> unit (** Used to end a local scope, purging the symbol table of all symbols added in that scope *) +val end_scope : 'a state -> unit -val set_read_only : 'a state -> string -> unit (** Used to add a read only label to an identifier *) +val set_read_only : 'a state -> string -> unit -val get_read_only : 'a state -> string -> bool (** Used to check for a read only label for an identifier *) +val get_read_only : 'a state -> string -> bool -val set_is_assigned : 'a state -> string -> unit (** Label an identifier as having been assigned to *) +val set_is_assigned : 'a state -> string -> unit -val set_is_unassigned : 'a state -> string -> unit (** Label an identifier as not having been assigned to *) +val set_is_unassigned : 'a state -> string -> unit -val check_is_unassigned : 'a state -> string -> bool (** Check whether an identifier is labelled as unassigned *) +val check_is_unassigned : 'a state -> string -> bool -val check_some_id_is_unassigned : 'a state -> bool (** Used to check whether some identifier is labelled as unassigned *) +val check_some_id_is_unassigned : 'a state -> bool -val is_global : 'a state -> string -> bool (** Used to check whether an identifier was declared in global scope *) +val is_global : 'a state -> string -> bool -val unsafe_clear_symbol_table : 'a state -> unit (** Used to clear the whole symbol table *) +val unsafe_clear_symbol_table : 'a state -> unit diff --git a/src/middle/Compiler.ml b/src/middle/Compiler.ml index 5c3d7966d9..261076d925 100644 --- a/src/middle/Compiler.ml +++ b/src/middle/Compiler.ml @@ -11,13 +11,13 @@ module type Frontend = sig val render_error : frontend_error -> string - val mir_of_file : - opts:frontend_opts + val mir_of_file + : opts:frontend_opts -> file:string -> (Program.Typed.t, frontend_error) result - val mir_of_string : - opts:frontend_opts + val mir_of_string + : opts:frontend_opts -> str:string -> (Program.Typed.t, frontend_error) result end @@ -40,9 +40,7 @@ module type Optimization = sig type optimization_opts (* parse level from string, for use in e.g. command line argument parser *) - val optimization_opts_of_string : - string -> (optimization_opts, string) result - + val optimization_opts_of_string : string -> (optimization_opts, string) result val default_optimization_opts : optimization_opts val optimize : opts:optimization_opts -> Program.Typed.t -> Program.Typed.t end @@ -55,24 +53,30 @@ module Compiler = struct val default_compiler_opts : compiler_opts - val compiler_opts_of_string : - string -> (compiler_opts, compiler_opts_error list) result + val compiler_opts_of_string + : string + -> (compiler_opts, compiler_opts_error list) result - val compile_from_file : - opts:compiler_opts -> file:string -> (string, frontend_error) result + val compile_from_file + : opts:compiler_opts + -> file:string + -> (string, frontend_error) result end module Make (F : Frontend) (O : Optimization) (B : Backend) : S with type frontend_error := F.frontend_error = struct type compiler_opts = - { frontend_opts: F.frontend_opts - ; optimization_opts: O.optimization_opts - ; backend_opts: B.backend_opts } + { frontend_opts : F.frontend_opts + ; optimization_opts : O.optimization_opts + ; backend_opts : B.backend_opts + } let default_compiler_opts = - { frontend_opts= F.default_frontend_opts - ; optimization_opts= O.default_optimization_opts - ; backend_opts= B.default_backend_opts } + { frontend_opts = F.default_frontend_opts + ; optimization_opts = O.default_optimization_opts + ; backend_opts = B.default_backend_opts + } + ;; [@@@ocaml.warning "-37"] @@ -81,12 +85,12 @@ module Compiler = struct | Optimize_opts_error of string | Backend_opts_error of string - let compiler_opts_of_string str = - Error [Frontend_opts_error ("todo " ^ str)] + let compiler_opts_of_string str = Error [ Frontend_opts_error ("todo " ^ str) ] let compile_from_file ~opts ~file = F.mir_of_file ~opts:opts.frontend_opts ~file |> Result.map ~f:(O.optimize ~opts:opts.optimization_opts) |> Result.map ~f:(B.mir_to_string ~opts:opts.backend_opts) + ;; end end diff --git a/src/middle/Expr.ml b/src/middle/Expr.ml index 21c13f54ff..29b61d284e 100644 --- a/src/middle/Expr.ml +++ b/src/middle/Expr.ml @@ -5,7 +5,11 @@ open Helpers (** Pattern and fixed-point of MIR expressions *) module Fixed = struct module Pattern = struct - type litType = Int | Real | Str [@@deriving sexp, hash, compare] + type litType = + | Int + | Real + | Str + [@@deriving sexp, hash, compare] type 'a t = | Var of string @@ -21,28 +25,52 @@ module Fixed = struct | Var varname -> Fmt.string ppf varname | Lit (Str, str) -> Fmt.pf ppf "%S" str | Lit (_, str) -> Fmt.string ppf str - | FunApp (StanLib, name, [lhs; rhs]) + | FunApp (StanLib, name, [ lhs; rhs ]) when Option.is_some (Operator.of_string_opt name) -> - Fmt.pf ppf "(%a %a %a)" pp_e lhs Operator.pp - (Option.value_exn (Operator.of_string_opt name)) - pp_e rhs + Fmt.pf + ppf + "(%a %a %a)" + pp_e + lhs + Operator.pp + (Option.value_exn (Operator.of_string_opt name)) + pp_e + rhs | FunApp (_, name, args) -> - Fmt.string ppf name ; - Fmt.(list pp_e ~sep:Fmt.comma |> parens) ppf args + Fmt.string ppf name; + Fmt.(list pp_e ~sep:Fmt.comma |> parens) ppf args | TernaryIf (pred, texpr, fexpr) -> - Fmt.pf ppf {|@[%a@ %a@,%a@,%a@ %a@]|} pp_e pred pp_builtin_syntax "?" - pp_e texpr pp_builtin_syntax ":" pp_e fexpr + Fmt.pf + ppf + {|@[%a@ %a@,%a@,%a@ %a@]|} + pp_e + pred + pp_builtin_syntax + "?" + pp_e + texpr + pp_builtin_syntax + ":" + pp_e + fexpr | Indexed (expr, indices) -> - Fmt.pf ppf {|@[%a%a@]|} pp_e expr - ( if List.is_empty indices then fun _ _ -> () - else Fmt.(list (Index.pp pp_e) ~sep:comma |> brackets) ) - indices + Fmt.pf + ppf + {|@[%a%a@]|} + pp_e + expr + (if List.is_empty indices + then fun _ _ -> () + else Fmt.(list (Index.pp pp_e) ~sep:comma |> brackets)) + indices | EAnd (l, r) -> Fmt.pf ppf "%a && %a" pp_e l pp_e r | EOr (l, r) -> Fmt.pf ppf "%a || %a" pp_e l pp_e r + ;; - include Foldable.Make (struct type nonrec 'a t = 'a t + include Foldable.Make (struct + type nonrec 'a t = 'a t - let fold = fold + let fold = fold end) end @@ -67,50 +95,59 @@ end module Typed = struct module Meta = struct type t = - { type_: UnsizedType.t - ; loc: Location_span.t sexp_opaque [@compare.ignore] - ; adlevel: UnsizedType.autodifftype } + { type_ : UnsizedType.t + ; loc : Location_span.t sexp_opaque [@compare.ignore] + ; adlevel : UnsizedType.autodifftype + } [@@deriving compare, create, sexp, hash] let empty = - create ~type_:UnsizedType.UInt ~adlevel:UnsizedType.DataOnly - ~loc:Location_span.empty () + create + ~type_:UnsizedType.UInt + ~adlevel:UnsizedType.DataOnly + ~loc:Location_span.empty + () + ;; let pp _ _ = () end include Specialized.Make (Fixed) (Meta) - let type_of Fixed.({meta= Meta.({type_; _}); _}) = type_ - let loc_of Fixed.({meta= Meta.({loc; _}); _}) = loc - let adlevel_of Fixed.({meta= Meta.({adlevel; _}); _}) = adlevel + let type_of Fixed.{ meta = Meta.{ type_; _ }; _ } = type_ + let loc_of Fixed.{ meta = Meta.{ loc; _ }; _ } = loc + let adlevel_of Fixed.{ meta = Meta.{ adlevel; _ }; _ } = adlevel end (** Expressions with associated location, type and label *) module Labelled = struct module Meta = struct type t = - { type_: UnsizedType.t - ; loc: Location_span.t sexp_opaque [@compare.ignore] - ; adlevel: UnsizedType.autodifftype - ; label: Label.Int_label.t [@compare.ignore] } + { type_ : UnsizedType.t + ; loc : Location_span.t sexp_opaque [@compare.ignore] + ; adlevel : UnsizedType.autodifftype + ; label : Label.Int_label.t [@compare.ignore] + } [@@deriving compare, create, sexp, hash] let empty = - create ~type_:UnsizedType.UInt ~adlevel:UnsizedType.DataOnly + create + ~type_:UnsizedType.UInt + ~adlevel:UnsizedType.DataOnly ~loc:Location_span.empty ~label:Label.Int_label.(prev init) () + ;; let pp _ _ = () end include Specialized.Make (Fixed) (Meta) - let type_of Fixed.({meta= Meta.({type_; _}); _}) = type_ - let label_of Fixed.({meta= Meta.({label; _}); _}) = label - let adlevel_of Fixed.({meta= Meta.({adlevel; _}); _}) = adlevel - let loc_of Fixed.({meta= Meta.({loc; _}); _}) = loc + let type_of Fixed.{ meta = Meta.{ type_; _ }; _ } = type_ + let label_of Fixed.{ meta = Meta.{ label; _ }; _ } = label + let adlevel_of Fixed.{ meta = Meta.{ adlevel; _ }; _ } = adlevel + let loc_of Fixed.{ meta = Meta.{ loc; _ }; _ } = loc (** Traverse a typed expression adding unique labels using locally mutable state @@ -118,62 +155,67 @@ module Labelled = struct let label ?(init = Label.Int_label.init) (expr : Typed.t) : t = let lbl = ref init in Fixed.map - (fun Typed.Meta.({adlevel; type_; loc}) -> + (fun Typed.Meta.{ adlevel; type_; loc } -> let cur_lbl = !lbl in - lbl := Label.Int_label.next cur_lbl ; - Meta.create ~label:cur_lbl ~adlevel ~type_ ~loc () ) + lbl := Label.Int_label.next cur_lbl; + Meta.create ~label:cur_lbl ~adlevel ~type_ ~loc ()) expr + ;; (** Build a map from expression labels to expressions *) - let rec associate ?init:(assocs = Label.Int_label.Map.empty) - ({pattern; _} as expr : t) = + let rec associate + ?init:(assocs = Label.Int_label.Map.empty) + ({ pattern; _ } as expr : t) + = let assocs_result : t Label.Int_label.Map.t Map_intf.Or_duplicate.t = - Label.Int_label.Map.add ~key:(label_of expr) ~data:expr + Label.Int_label.Map.add + ~key:(label_of expr) + ~data:expr (associate_pattern assocs @@ pattern) in - match assocs_result with `Ok x -> x | `Duplicate -> assocs + match assocs_result with + | `Ok x -> x + | `Duplicate -> assocs and associate_pattern assocs = function | Fixed.Pattern.Lit _ | Var _ -> assocs | FunApp (_, _, args) -> - List.fold args ~init:assocs ~f:(fun accu x -> associate ~init:accu x) - | EAnd (e1, e2) | EOr (e1, e2) -> - associate ~init:(associate ~init:assocs e2) e1 + List.fold args ~init:assocs ~f:(fun accu x -> associate ~init:accu x) + | EAnd (e1, e2) | EOr (e1, e2) -> associate ~init:(associate ~init:assocs e2) e1 | TernaryIf (e1, e2, e3) -> - associate ~init:(associate ~init:(associate ~init:assocs e3) e2) e1 + associate ~init:(associate ~init:(associate ~init:assocs e3) e2) e1 | Indexed (e, idxs) -> - List.fold idxs ~init:(associate ~init:assocs e) ~f:associate_index + List.fold idxs ~init:(associate ~init:assocs e) ~f:associate_index and associate_index assocs = function | All -> assocs | Single e | Upfrom e | MultiIndex e -> associate ~init:assocs e | Between (e1, e2) -> associate ~init:(associate ~init:assocs e2) e1 + ;; end module Helpers = struct - let int i = - {Fixed.meta= Typed.Meta.empty; pattern= Lit (Int, string_of_int i)} - - let float i = - {Fixed.meta= Typed.Meta.empty; pattern= Lit (Real, string_of_float i)} - - let str i = {Fixed.meta= Typed.Meta.empty; pattern= Lit (Str, i)} + let int i = { Fixed.meta = Typed.Meta.empty; pattern = Lit (Int, string_of_int i) } + let float i = { Fixed.meta = Typed.Meta.empty; pattern = Lit (Real, string_of_float i) } + let str i = { Fixed.meta = Typed.Meta.empty; pattern = Lit (Str, i) } let zero = int 0 let one = int 1 let binop e1 op e2 = - { Fixed.meta= Typed.Meta.empty - ; pattern= FunApp (StanLib, Operator.to_string op, [e1; e2]) } + { Fixed.meta = Typed.Meta.empty + ; pattern = FunApp (StanLib, Operator.to_string op, [ e1; e2 ]) + } + ;; let loop_bottom = one let internal_funapp fn args meta = - { Fixed.meta - ; pattern= FunApp (CompilerInternal, Internal_fun.to_string fn, args) } + { Fixed.meta; pattern = FunApp (CompilerInternal, Internal_fun.to_string fn, args) } + ;; let contains_fn fn ?(init = false) e = let fstr = Internal_fun.to_string fn in - let rec aux accu Fixed.({pattern; _}) = + let rec aux accu Fixed.{ pattern; _ } = accu || match pattern with @@ -181,57 +223,58 @@ module Helpers = struct | x -> Fixed.Pattern.fold aux accu x in aux init e + ;; - let%test "expr contains fn" = - internal_funapp FnReadData [] () |> contains_fn FnReadData + let%test "expr contains fn" = internal_funapp FnReadData [] () |> contains_fn FnReadData let rec infer_type_of_indexed ut indices = - match (ut, indices) with + match ut, indices with | _, [] -> ut - | _, [Index.All] | _, [Upfrom _] | _, [Between _] -> ut - | UnsizedType.UMatrix, [All; Single _] - |UMatrix, [Upfrom _; Single _] - |UMatrix, [Between _; Single _] - |UMatrix, [MultiIndex _] - |UMatrix, [Single _] -> - UVector + | _, [ Index.All ] | _, [ Upfrom _ ] | _, [ Between _ ] -> ut + | UnsizedType.UMatrix, [ All; Single _ ] + | UMatrix, [ Upfrom _; Single _ ] + | UMatrix, [ Between _; Single _ ] + | UMatrix, [ MultiIndex _ ] + | UMatrix, [ Single _ ] -> UVector | UArray t, Single _ :: tl -> infer_type_of_indexed t tl | UArray t, _ :: tl -> UArray (infer_type_of_indexed t tl) - | UMatrix, [Single _; Single _] | UVector, [_] | URowVector, [_] -> UReal + | UMatrix, [ Single _; Single _ ] | UVector, [ _ ] | URowVector, [ _ ] -> UReal | _ -> raise_s [%message "Can't index" (ut : UnsizedType.t)] + ;; (** [add_index expression index] returns an expression that (additionally) indexes into the input [expression] by [index].*) let add_int_index e i = - let mtype = infer_type_of_indexed Typed.(type_of e) [i] in - let meta = Typed.Meta.{e.meta with type_= mtype} + let mtype = infer_type_of_indexed Typed.(type_of e) [ i ] in + let meta = Typed.Meta.{ e.meta with type_ = mtype } and pattern = match e.pattern with - | Var _ -> Fixed.Pattern.Indexed (e, [i]) - | Indexed (e, indices) -> Indexed (e, indices @ [i]) + | Var _ -> Fixed.Pattern.Indexed (e, [ i ]) + | Indexed (e, indices) -> Indexed (e, indices @ [ i ]) | _ -> raise_s [%message "These should go away with Ryan's LHS"] in - Fixed.{meta; pattern} + Fixed.{ meta; pattern } + ;; (** TODO: Make me tail recursive *) - let rec collect_indices Fixed.({pattern; _}) = + let rec collect_indices Fixed.{ pattern; _ } = match pattern with | Indexed (obj, indices) -> collect_indices obj @ indices | _ -> [] + ;; let%expect_test "infer type of indexed" = - [ ( UnsizedType.UArray UMatrix - , [Index.Single loop_bottom; Single loop_bottom] ) - ; (UArray (UArray UMatrix), [Single loop_bottom]) - ; (UArray UMatrix, [Single loop_bottom]) - ; (UArray UMatrix, [Upfrom loop_bottom; Single loop_bottom]) - ; ( UArray UMatrix - , [Single loop_bottom; Single loop_bottom; Single loop_bottom] ) - ; ( UArray UMatrix - , [Upfrom loop_bottom; Single loop_bottom; Single loop_bottom] ) ] + [ UnsizedType.UArray UMatrix, [ Index.Single loop_bottom; Single loop_bottom ] + ; UArray (UArray UMatrix), [ Single loop_bottom ] + ; UArray UMatrix, [ Single loop_bottom ] + ; UArray UMatrix, [ Upfrom loop_bottom; Single loop_bottom ] + ; UArray UMatrix, [ Single loop_bottom; Single loop_bottom; Single loop_bottom ] + ; UArray UMatrix, [ Upfrom loop_bottom; Single loop_bottom; Single loop_bottom ] + ] |> List.map ~f:(fun (ut, idx) -> infer_type_of_indexed ut idx) |> Fmt.(strf "@[%a@]" (list ~sep:comma UnsizedType.pp)) - |> print_endline ; + |> print_endline; [%expect {| vector, matrix[], matrix, vector[], real, real[] |}] + ;; end diff --git a/src/middle/Expr.mli b/src/middle/Expr.mli index 5d8879a482..463051c319 100644 --- a/src/middle/Expr.mli +++ b/src/middle/Expr.mli @@ -3,7 +3,11 @@ open Common module Fixed : sig module Pattern : sig - type litType = Int | Real | Str [@@deriving sexp, hash, compare] + type litType = + | Int + | Real + | Str + [@@deriving sexp, hash, compare] type 'a t = | Var of string @@ -36,9 +40,10 @@ end module Typed : sig module Meta : sig type t = - { type_: UnsizedType.t - ; loc: Location_span.t sexp_opaque [@compare.ignore] - ; adlevel: UnsizedType.autodifftype } + { type_ : UnsizedType.t + ; loc : Location_span.t sexp_opaque [@compare.ignore] + ; adlevel : UnsizedType.autodifftype + } [@@deriving compare, create, sexp, hash] include Specialized.Meta with type t := t @@ -54,10 +59,11 @@ end module Labelled : sig module Meta : sig type t = - { type_: UnsizedType.t - ; loc: Location_span.t sexp_opaque [@compare.ignore] - ; adlevel: UnsizedType.autodifftype - ; label: Label.Int_label.t } + { type_ : UnsizedType.t + ; loc : Location_span.t sexp_opaque [@compare.ignore] + ; adlevel : UnsizedType.autodifftype + ; label : Label.Int_label.t + } [@@deriving compare, create, sexp, hash] include Specialized.Meta with type t := t @@ -71,9 +77,7 @@ module Labelled : sig val label_of : t -> Label.Int_label.t val label : ?init:int -> Typed.t -> t val associate : ?init:t Label.Int_label.Map.t -> t -> t Label.Int_label.Map.t - - val associate_index : - t Label.Int_label.Map.t -> t Index.t -> t Label.Int_label.Map.t + val associate_index : t Label.Int_label.Map.t -> t Index.t -> t Label.Int_label.Map.t end module Helpers : sig diff --git a/src/middle/Flag_vars.ml b/src/middle/Flag_vars.ml index fd3b2f69a9..f9aa63eda8 100644 --- a/src/middle/Flag_vars.ml +++ b/src/middle/Flag_vars.ml @@ -1,7 +1,10 @@ -type t = EmitGeneratedQuantities | EmitTransformedParameters +type t = + | EmitGeneratedQuantities + | EmitTransformedParameters -let enumerate = [EmitGeneratedQuantities; EmitTransformedParameters] +let enumerate = [ EmitGeneratedQuantities; EmitTransformedParameters ] let to_string = function | EmitGeneratedQuantities -> "emit_generated_quantities__" | EmitTransformedParameters -> "emit_transformed_parameters__" +;; diff --git a/src/middle/Fun_kind.ml b/src/middle/Fun_kind.ml index 7728311a6a..d44e9e3bfa 100644 --- a/src/middle/Fun_kind.ml +++ b/src/middle/Fun_kind.ml @@ -1,2 +1,5 @@ -type t = StanLib | CompilerInternal | UserDefined +type t = + | StanLib + | CompilerInternal + | UserDefined [@@deriving compare, sexp, hash] diff --git a/src/middle/Index.ml b/src/middle/Index.ml index 080e97f71d..6c3340a635 100644 --- a/src/middle/Index.ml +++ b/src/middle/Index.ml @@ -14,14 +14,21 @@ let pp pp_e ppf = function | Upfrom index -> Fmt.pf ppf {|%a:|} pp_e index | Between (lower, upper) -> Fmt.pf ppf {|%a:%a|} pp_e lower pp_e upper | MultiIndex index -> Fmt.pf ppf {|%a|} pp_e index +;; let pp_indexed pp_e ppf (ident, indices) = - Fmt.pf ppf {|@[%s%a@]|} ident - ( if List.is_empty indices then fun _ _ -> () - else Fmt.(list (pp pp_e) ~sep:comma |> brackets) ) + Fmt.pf + ppf + {|@[%s%a@]|} + ident + (if List.is_empty indices + then fun _ _ -> () + else Fmt.(list (pp pp_e) ~sep:comma |> brackets)) indices +;; let bounds = function | All -> [] - | Single e | Upfrom e | MultiIndex e -> [e] - | Between (e1, e2) -> [e1; e2] + | Single e | Upfrom e | MultiIndex e -> [ e ] + | Between (e1, e2) -> [ e1; e2 ] +;; diff --git a/src/middle/Internal_fun.ml b/src/middle/Internal_fun.ml index e889bedf49..72657243af 100644 --- a/src/middle/Internal_fun.ml +++ b/src/middle/Internal_fun.ml @@ -24,9 +24,7 @@ type t = let to_string x = Sexp.to_string (sexp_of_t x) ^ "__" let of_string_opt x = - try - String.chop_suffix_exn ~suffix:"__" x - |> Sexp.of_string |> t_of_sexp |> Some - with + try String.chop_suffix_exn ~suffix:"__" x |> Sexp.of_string |> t_of_sexp |> Some with | Sexplib.Conv.Of_sexp_error _ -> None | Invalid_argument _ -> None +;; diff --git a/src/middle/Location.ml b/src/middle/Location.ml index b79cf695e5..77b53a6740 100644 --- a/src/middle/Location.ml +++ b/src/middle/Location.ml @@ -3,20 +3,25 @@ module Str = Re.Str (** Source code locations *) type t = - {filename: string; line_num: int; col_num: int; included_from: t option} + { filename : string + ; line_num : int + ; col_num : int + ; included_from : t option + } [@@deriving sexp, hash, compare] -let pp_context_exn ppf {filename; line_num; col_num; _} = +let pp_context_exn ppf { filename; line_num; col_num; _ } = let open In_channel in let input = create filename in for _ = 1 to line_num - 3 do ignore (input_line_exn input) - done ; + done; let get_line num = - if num > 0 then + if num > 0 + then ( match input_line input with | Some input -> Printf.sprintf "%6d: %s\n" num input - | _ -> "" + | _ -> "") else "" in let line_2_before = get_line (line_num - 2) in @@ -25,21 +30,31 @@ let pp_context_exn ppf {filename; line_num; col_num; _} = let cursor_line = String.make (col_num + 9) ' ' ^ "^\n" in let line_after = get_line (line_num + 1) in let line_2_after = get_line (line_num + 2) in - close input ; - Fmt.pf ppf + close input; + Fmt.pf + ppf " -------------------------------------------------\n\ %s%s%s%s%s%s -------------------------------------------------\n" - line_2_before line_before our_line cursor_line line_after line_2_after + line_2_before + line_before + our_line + cursor_line + line_after + line_2_after +;; let context_to_string file = - try Some (Fmt.to_to_string pp_context_exn file) with _ -> None + try Some (Fmt.to_to_string pp_context_exn file) with + | _ -> None +;; (** Return two lines before and after the specified location and print a message *) let pp_with_message_exn ppf (message, loc) = Fmt.pf ppf "%a\n%s\n\n" pp_context_exn loc message +;; -let empty = {filename= ""; line_num= 0; col_num= 0; included_from= None} +let empty = { filename = ""; line_num = 0; col_num = 0; included_from = None } let rec to_string ?(print_file = true) ?(print_line = true) loc = let open Format in @@ -51,51 +66,55 @@ let rec to_string ?(print_file = true) ?(print_line = true) loc = | None -> "" in sprintf "%s%scolumn %d%s" file line loc.col_num incl +;; let trim_quotes s = let s = String.drop_prefix s 1 in String.drop_suffix s 1 +;; let rec of_string_opt str = let split_str = - Str.bounded_split - (Str.regexp ", line \\|, column \\|, included from\n") - str 4 + Str.bounded_split (Str.regexp ", line \\|, column \\|, included from\n") str 4 in match split_str with - | [fname; linenum_str; colnum_str] -> - Some - { filename= trim_quotes fname - ; line_num= int_of_string linenum_str - ; col_num= int_of_string colnum_str - ; included_from= None } - | [fname; linenum_str; colnum_str; included_from_str] -> - of_string_opt included_from_str - |> Option.map ~f:(fun included_from -> - { filename= trim_quotes fname - ; line_num= int_of_string linenum_str - ; col_num= int_of_string colnum_str - ; included_from= Some included_from } ) + | [ fname; linenum_str; colnum_str ] -> + Some + { filename = trim_quotes fname + ; line_num = int_of_string linenum_str + ; col_num = int_of_string colnum_str + ; included_from = None + } + | [ fname; linenum_str; colnum_str; included_from_str ] -> + of_string_opt included_from_str + |> Option.map ~f:(fun included_from -> + { filename = trim_quotes fname + ; line_num = int_of_string linenum_str + ; col_num = int_of_string colnum_str + ; included_from = Some included_from + }) | _ -> None +;; -let of_position_opt {Lexing.pos_fname; pos_lnum; pos_cnum; pos_bol} = - let split_fname = - Str.bounded_split (Str.regexp ", included from\n") pos_fname 2 - in +let of_position_opt { Lexing.pos_fname; pos_lnum; pos_cnum; pos_bol } = + let split_fname = Str.bounded_split (Str.regexp ", included from\n") pos_fname 2 in match split_fname with | [] -> None - | [fname] -> - Some - { filename= fname - ; line_num= pos_lnum - ; col_num= pos_cnum - pos_bol - ; included_from= None } + | [ fname ] -> + Some + { filename = fname + ; line_num = pos_lnum + ; col_num = pos_cnum - pos_bol + ; included_from = None + } | fname1 :: fname2 :: _ -> - Option.map (of_string_opt fname2) ~f:(fun included_from -> - { filename= fname1 - ; line_num= pos_lnum - ; col_num= pos_cnum - pos_bol - ; included_from= Some included_from } ) + Option.map (of_string_opt fname2) ~f:(fun included_from -> + { filename = fname1 + ; line_num = pos_lnum + ; col_num = pos_cnum - pos_bol + ; included_from = Some included_from + }) +;; let of_position_exn fn = Option.value_exn (of_position_opt fn) @@ -106,39 +125,38 @@ let%expect_test "location string equivalence 1" = 'yyy.stan', line 666, column 42, included from\n\ 'zzz.stan', line 24, column 77" in - print_endline (to_string @@ Option.value_exn (of_string_opt str)) ; + print_endline (to_string @@ Option.value_exn (of_string_opt str)); [%expect {| 'xxx.stan', line 245, column 13, included from 'yyy.stan', line 666, column 42, included from 'zzz.stan', line 24, column 77 |}] +;; let%expect_test "location string equivalence 2" = let loc : t = - { filename= "xxx.stan" - ; line_num= 35 - ; col_num= 24 - ; included_from= + { filename = "xxx.stan" + ; line_num = 35 + ; col_num = 24 + ; included_from = Some - { filename= "yyy.stan" - ; line_num= 345 - ; col_num= 214 - ; included_from= None } } + { filename = "yyy.stan"; line_num = 345; col_num = 214; included_from = None } + } in - print_endline (to_string @@ Option.value_exn (of_string_opt (to_string loc))) ; + print_endline (to_string @@ Option.value_exn (of_string_opt (to_string loc))); [%expect {| 'xxx.stan', line 35, column 24, included from 'yyy.stan', line 345, column 214 |}] +;; let%expect_test "parse location from string" = - let loc = - Option.value_exn (of_string_opt "'xxx.stan', line 245, column 13") - in - print_endline loc.filename ; - print_endline (string_of_int loc.line_num) ; - print_endline (string_of_int loc.col_num) ; + let loc = Option.value_exn (of_string_opt "'xxx.stan', line 245, column 13") in + print_endline loc.filename; + print_endline (string_of_int loc.line_num); + print_endline (string_of_int loc.col_num); [%expect {| xxx.stan 245 13 |}] +;; diff --git a/src/middle/Location_span.ml b/src/middle/Location_span.ml index c2958a648d..90908bfaae 100644 --- a/src/middle/Location_span.ml +++ b/src/middle/Location_span.ml @@ -2,36 +2,41 @@ open Core_kernel (** Delimited locations *) -type t = {begin_loc: Location.t; end_loc: Location.t} +type t = + { begin_loc : Location.t + ; end_loc : Location.t + } [@@deriving sexp, hash, compare] -let empty = {begin_loc= Location.empty; end_loc= Location.empty} -let merge left right = {begin_loc= left.begin_loc; end_loc= right.end_loc} +let empty = { begin_loc = Location.empty; end_loc = Location.empty } +let merge left right = { begin_loc = left.begin_loc; end_loc = right.end_loc } (** Render a location_span as a string *) -let to_string {begin_loc; end_loc} = +let to_string { begin_loc; end_loc } = let end_loc_str = match begin_loc.included_from with | None -> - " to " - ^ Location.to_string - ~print_file:(begin_loc.filename <> end_loc.filename) - ~print_line:(begin_loc.line_num <> end_loc.line_num) - end_loc + " to " + ^ Location.to_string + ~print_file:(begin_loc.filename <> end_loc.filename) + ~print_line:(begin_loc.line_num <> end_loc.line_num) + end_loc | Some _ -> "" in Location.to_string begin_loc ^ end_loc_str +;; (** Take the Middle.location_span corresponding to a pair of Lexing.position's *) let of_positions_opt start_pos end_pos = Option.( Location.of_position_opt start_pos >>= fun begin_loc -> - Location.of_position_opt end_pos - |> map ~f:(fun end_loc -> {begin_loc; end_loc})) + Location.of_position_opt end_pos |> map ~f:(fun end_loc -> { begin_loc; end_loc })) +;; let of_positions_exn start_pos end_pos = Option.value_exn (of_positions_opt start_pos end_pos) +;; module Comparator = Comparator.Make (struct type nonrec t = t diff --git a/src/middle/Operator.ml b/src/middle/Operator.ml index 84aff36a64..399f7f46d9 100644 --- a/src/middle/Operator.ml +++ b/src/middle/Operator.ml @@ -46,16 +46,15 @@ let pp ppf = function | Geq -> Fmt.pf ppf ">=" | PNot -> Fmt.pf ppf "!" | Transpose -> Fmt.pf ppf "'" +;; let to_string x = Sexp.to_string (sexp_of_t x) ^ "__" let of_string_opt x = - try - String.chop_suffix_exn ~suffix:"__" x - |> Sexp.of_string |> t_of_sexp |> Some - with + try String.chop_suffix_exn ~suffix:"__" x |> Sexp.of_string |> t_of_sexp |> Some with | Sexplib.Conv.Of_sexp_error _ -> None | Invalid_argument _ -> None +;; let stan_math_name = function | Plus -> "add" @@ -64,8 +63,8 @@ let stan_math_name = function | PMinus -> "minus" | Times -> "multiply" (* TODO: this was taken from `Mir_utils.string_of_operators` - what was the intended behaviour here? - *) + what was the intended behaviour here? + *) | Divide -> "mdivide_right" (* | Divide -> "divide" *) | IntDivide -> "divide" @@ -84,3 +83,4 @@ let stan_math_name = function | Geq -> "logical_gte" | PNot -> "logical_negation" | Transpose -> "transpose" +;; diff --git a/src/middle/Program.ml b/src/middle/Program.ml index 5b93470ff4..89f8e8d1c7 100644 --- a/src/middle/Program.ml +++ b/src/middle/Program.ml @@ -6,14 +6,18 @@ type fun_arg_decl = (UnsizedType.autodifftype * string * UnsizedType.t) list [@@deriving sexp, hash, map] type 'a fun_def = - { fdrt: UnsizedType.t option - ; fdname: string - ; fdargs: (UnsizedType.autodifftype * string * UnsizedType.t) list - ; fdbody: 'a - ; fdloc: Location_span.t sexp_opaque [@compare.ignore] } + { fdrt : UnsizedType.t option + ; fdname : string + ; fdargs : (UnsizedType.autodifftype * string * UnsizedType.t) list + ; fdbody : 'a + ; fdloc : Location_span.t sexp_opaque [@compare.ignore] + } [@@deriving compare, hash, map, sexp, map, fold] -type io_block = Parameters | TransformedParameters | GeneratedQuantities +type io_block = + | Parameters + | TransformedParameters + | GeneratedQuantities [@@deriving sexp, hash] (** Transformations (constraints) for global variable declarations *) @@ -36,105 +40,150 @@ type 'e transformation = [@@deriving sexp, compare, map, hash, fold] type 'e outvar = - { out_unconstrained_st: 'e SizedType.t - ; out_constrained_st: 'e SizedType.t - ; out_block: io_block - ; out_trans: 'e transformation } + { out_unconstrained_st : 'e SizedType.t + ; out_constrained_st : 'e SizedType.t + ; out_block : io_block + ; out_trans : 'e transformation + } [@@deriving sexp, map, hash, fold] type ('a, 'b) t = - { functions_block: 'b fun_def list - ; input_vars: (string * 'a SizedType.t) list - ; prepare_data: 'b list (* data & transformed data decls and statements *) - ; log_prob: 'b list (*assumes data & params are in scope and ready*) - ; generate_quantities: 'b list (* assumes data & params ready & in scope*) - ; transform_inits: 'b list - ; output_vars: (string * 'a outvar) list - ; prog_name: string - ; prog_path: string } + { functions_block : 'b fun_def list + ; input_vars : (string * 'a SizedType.t) list + ; prepare_data : 'b list (* data & transformed data decls and statements *) + ; log_prob : 'b list (*assumes data & params are in scope and ready*) + ; generate_quantities : 'b list (* assumes data & params ready & in scope*) + ; transform_inits : 'b list + ; output_vars : (string * 'a outvar) list + ; prog_name : string + ; prog_path : string + } [@@deriving sexp, map, fold] let map_stmts f p = { p with - prepare_data= f p.prepare_data - ; log_prob= f p.log_prob - ; generate_quantities= f p.generate_quantities - ; transform_inits= f p.transform_inits } + prepare_data = f p.prepare_data + ; log_prob = f p.log_prob + ; generate_quantities = f p.generate_quantities + ; transform_inits = f p.transform_inits + } +;; (* -- Pretty printers -- *) let pp_fun_arg_decl ppf (autodifftype, name, unsizedtype) = - Fmt.pf ppf "%a%a %s" UnsizedType.pp_autodifftype autodifftype UnsizedType.pp - unsizedtype name + Fmt.pf + ppf + "%a%a %s" + UnsizedType.pp_autodifftype + autodifftype + UnsizedType.pp + unsizedtype + name +;; let pp_fun_def pp_s ppf = function - | {fdrt; fdname; fdargs; fdbody; _} -> ( - match fdrt with + | { fdrt; fdname; fdargs; fdbody; _ } -> + (match fdrt with | Some rt -> - Fmt.pf ppf {|@[%a %s%a {@ %a@]@ }|} UnsizedType.pp rt fdname - Fmt.(list pp_fun_arg_decl ~sep:comma |> parens) - fdargs pp_s fdbody + Fmt.pf + ppf + {|@[%a %s%a {@ %a@]@ }|} + UnsizedType.pp + rt + fdname + Fmt.(list pp_fun_arg_decl ~sep:comma |> parens) + fdargs + pp_s + fdbody | None -> - Fmt.pf ppf {|@[%s %s%a {@ %a@]@ }|} "void" fdname - Fmt.(list pp_fun_arg_decl ~sep:comma |> parens) - fdargs pp_s fdbody ) + Fmt.pf + ppf + {|@[%s %s%a {@ %a@]@ }|} + "void" + fdname + Fmt.(list pp_fun_arg_decl ~sep:comma |> parens) + fdargs + pp_s + fdbody) +;; let pp_io_block ppf = function | Parameters -> Fmt.string ppf "parameters" | TransformedParameters -> Fmt.string ppf "transformed_parameters" | GeneratedQuantities -> Fmt.string ppf "generated_quantities" +;; let pp_block label pp_elem ppf = function | [] -> () | elems -> - Fmt.pf ppf {|@[%a {@ %a@]@ }|} pp_keyword label - Fmt.(list ~sep:cut pp_elem) - elems ; - Format.pp_force_newline ppf () + Fmt.pf ppf {|@[%a {@ %a@]@ }|} pp_keyword label Fmt.(list ~sep:cut pp_elem) elems; + Format.pp_force_newline ppf () +;; -let pp_functions_block pp_s ppf {functions_block; _} = +let pp_functions_block pp_s ppf { functions_block; _ } = pp_block "functions" pp_s ppf functions_block +;; -let pp_prepare_data pp_s ppf {prepare_data; _} = +let pp_prepare_data pp_s ppf { prepare_data; _ } = pp_block "prepare_data" pp_s ppf prepare_data +;; -let pp_log_prob pp_s ppf {log_prob; _} = pp_block "log_prob" pp_s ppf log_prob +let pp_log_prob pp_s ppf { log_prob; _ } = pp_block "log_prob" pp_s ppf log_prob -let pp_generate_quantities pp_s ppf {generate_quantities; _} = +let pp_generate_quantities pp_s ppf { generate_quantities; _ } = pp_block "generate_quantities" pp_s ppf generate_quantities +;; -let pp_transform_inits pp_s ppf {transform_inits; _} = +let pp_transform_inits pp_s ppf { transform_inits; _ } = pp_block "transform_inits" pp_s ppf transform_inits - -let pp_output_var pp_e ppf - (name, {out_unconstrained_st; out_constrained_st; out_block; _}) = - Fmt.pf ppf "@[%a %a %s; //%a@]" pp_io_block out_block (SizedType.pp pp_e) - out_constrained_st name (SizedType.pp pp_e) out_unconstrained_st +;; + +let pp_output_var + pp_e + ppf + (name, { out_unconstrained_st; out_constrained_st; out_block; _ }) + = + Fmt.pf + ppf + "@[%a %a %s; //%a@]" + pp_io_block + out_block + (SizedType.pp pp_e) + out_constrained_st + name + (SizedType.pp pp_e) + out_unconstrained_st +;; let pp_input_var pp_e ppf (name, sized_ty) = Fmt.pf ppf "@[%a %s;@]" (SizedType.pp pp_e) sized_ty name +;; -let pp_input_vars pp_e ppf {input_vars; _} = +let pp_input_vars pp_e ppf { input_vars; _ } = pp_block "input_vars" (pp_input_var pp_e) ppf input_vars +;; -let pp_output_vars pp_e ppf {output_vars; _} = +let pp_output_vars pp_e ppf { output_vars; _ } = pp_block "output_vars" (pp_output_var pp_e) ppf output_vars +;; let pp pp_e pp_s ppf prog = - Format.open_vbox 0 ; - pp_functions_block (pp_fun_def pp_s) ppf prog ; - Fmt.cut ppf () ; - pp_input_vars pp_e ppf prog ; - Fmt.cut ppf () ; - pp_prepare_data pp_s ppf prog ; - Fmt.cut ppf () ; - pp_log_prob pp_s ppf prog ; - Fmt.cut ppf () ; - pp_generate_quantities pp_s ppf prog ; - Fmt.cut ppf () ; - pp_transform_inits pp_s ppf prog ; - Fmt.cut ppf () ; - pp_output_vars pp_e ppf prog ; + Format.open_vbox 0; + pp_functions_block (pp_fun_def pp_s) ppf prog; + Fmt.cut ppf (); + pp_input_vars pp_e ppf prog; + Fmt.cut ppf (); + pp_prepare_data pp_s ppf prog; + Fmt.cut ppf (); + pp_log_prob pp_s ppf prog; + Fmt.cut ppf (); + pp_generate_quantities pp_s ppf prog; + Fmt.cut ppf (); + pp_transform_inits pp_s ppf prog; + Fmt.cut ppf (); + pp_output_vars pp_e ppf prog; Format.close_box () +;; (* Programs with typed expressions and locations *) module Typed = struct @@ -154,66 +203,65 @@ module Labelled = struct let t_of_sexp = t_of_sexp Expr.Labelled.t_of_sexp Stmt.Labelled.t_of_sexp (* let label ?(init = 0) (prog : Typed.t) : t = - let incr_label = - State.(get >>= fun label -> put (label + 1) >>= fun _ -> return label) - in - let f {Expr.Typed.Meta.adlevel; type_; loc} = - incr_label - |> State.map ~f:(fun label -> - Expr.Labelled.Meta.create ~type_ ~loc ~adlevel ~label () ) - and g loc = - incr_label - |> State.map ~f:(fun label -> Stmt.Labelled.Meta.create ~loc ~label ()) - in - Traversable_state.traverse prog - ~f:(Traversable_expr_state.traverse ~f) - ~g:(Traversable_stmt_state.traverse ~f ~g) - |> State.run_state ~init |> fst *) + let incr_label = + State.(get >>= fun label -> put (label + 1) >>= fun _ -> return label) + in + let f {Expr.Typed.Meta.adlevel; type_; loc} = + incr_label + |> State.map ~f:(fun label -> + Expr.Labelled.Meta.create ~type_ ~loc ~adlevel ~label () ) + and g loc = + incr_label + |> State.map ~f:(fun label -> Stmt.Labelled.Meta.create ~loc ~label ()) + in + Traversable_state.traverse prog + ~f:(Traversable_expr_state.traverse ~f) + ~g:(Traversable_stmt_state.traverse ~f ~g) + |> State.run_state ~init |> fst *) let empty = - { Stmt.Labelled.exprs= Label.Int_label.Map.empty - ; stmts= Label.Int_label.Map.empty } + { Stmt.Labelled.exprs = Label.Int_label.Map.empty; stmts = Label.Int_label.Map.empty } + ;; let rec associate ?init:(assocs = empty) prog = let assoc_fundef = List.fold_left prog.functions_block ~init:assocs ~f:associate_fun_def in let assoc_input_vars = - List.fold_left prog.input_vars ~init:assoc_fundef - ~f:(fun assocs (_, st) -> - {assocs with exprs= SizedType.associate ~init:assocs.exprs st} ) + List.fold_left prog.input_vars ~init:assoc_fundef ~f:(fun assocs (_, st) -> + { assocs with exprs = SizedType.associate ~init:assocs.exprs st }) in let assoc_prepare_data = - List.fold_left prog.prepare_data ~init:assoc_input_vars - ~f:(fun assocs stmt -> Stmt.Labelled.associate ~init:assocs stmt ) + List.fold_left prog.prepare_data ~init:assoc_input_vars ~f:(fun assocs stmt -> + Stmt.Labelled.associate ~init:assocs stmt) in let assoc_log_prog = - List.fold_left prog.log_prob ~init:assoc_prepare_data - ~f:(fun assocs stmt -> Stmt.Labelled.associate ~init:assocs stmt ) + List.fold_left prog.log_prob ~init:assoc_prepare_data ~f:(fun assocs stmt -> + Stmt.Labelled.associate ~init:assocs stmt) in let assoc_generate_quants = - List.fold_left prog.generate_quantities ~init:assoc_log_prog - ~f:(fun assocs stmt -> Stmt.Labelled.associate ~init:assocs stmt ) + List.fold_left prog.generate_quantities ~init:assoc_log_prog ~f:(fun assocs stmt -> + Stmt.Labelled.associate ~init:assocs stmt) in let assoc_transform_inits = - List.fold_left prog.transform_inits ~init:assoc_generate_quants - ~f:(fun assocs stmt -> Stmt.Labelled.associate ~init:assocs stmt ) + List.fold_left + prog.transform_inits + ~init:assoc_generate_quants + ~f:(fun assocs stmt -> Stmt.Labelled.associate ~init:assocs stmt) in - List.fold_left prog.output_vars ~init:assoc_transform_inits - ~f:associate_outvar + List.fold_left prog.output_vars ~init:assoc_transform_inits ~f:associate_outvar - and associate_fun_def assocs {fdbody; _} = - Stmt.Labelled.associate ~init:assocs fdbody + and associate_fun_def assocs { fdbody; _ } = Stmt.Labelled.associate ~init:assocs fdbody - and associate_outvar assocs (_, {out_constrained_st; out_unconstrained_st; _}) - = + and associate_outvar assocs (_, { out_constrained_st; out_unconstrained_st; _ }) = let exprs = SizedType.( associate ~init:(associate ~init:assocs.exprs out_unconstrained_st) out_constrained_st) in - {assocs with exprs} + { assocs with exprs } + ;; end module Numbered = struct diff --git a/src/middle/SizedType.ml b/src/middle/SizedType.ml index 2f09832b25..36802f437d 100644 --- a/src/middle/SizedType.ml +++ b/src/middle/SizedType.ml @@ -16,13 +16,14 @@ let rec pp pp_e ppf = function | SVector expr -> Fmt.pf ppf {|vector%a|} (Fmt.brackets pp_e) expr | SRowVector expr -> Fmt.pf ppf {|row_vector%a|} (Fmt.brackets pp_e) expr | SMatrix (d1_expr, d2_expr) -> - Fmt.pf ppf {|matrix%a|} - Fmt.(pair ~sep:comma pp_e pp_e |> brackets) - (d1_expr, d2_expr) + Fmt.pf ppf {|matrix%a|} Fmt.(pair ~sep:comma pp_e pp_e |> brackets) (d1_expr, d2_expr) | SArray (st, expr) -> - Fmt.pf ppf {|array%a|} - Fmt.(pair ~sep:comma (fun ppf st -> pp pp_e ppf st) pp_e |> brackets) - (st, expr) + Fmt.pf + ppf + {|array%a|} + Fmt.(pair ~sep:comma (fun ppf st -> pp pp_e ppf st) pp_e |> brackets) + (st, expr) +;; let collect_exprs st = let rec aux accu = function @@ -32,6 +33,7 @@ let collect_exprs st = | SArray (inner, e) -> aux (e :: accu) inner in aux [] st +;; let rec to_unsized = function | SInt -> UnsizedType.UInt @@ -40,34 +42,43 @@ let rec to_unsized = function | SRowVector _ -> URowVector | SMatrix _ -> UMatrix | SArray (t, _) -> UArray (to_unsized t) +;; let rec associate ?init:(assocs = Label.Int_label.Map.empty) = function | SInt | SReal -> assocs | SVector e | SRowVector e -> Expr.Labelled.associate ~init:assocs e - | SMatrix (e1, e2) -> - Expr.Labelled.(associate ~init:(associate ~init:assocs e1) e2) - | SArray (st, e) -> - associate ~init:(Expr.Labelled.associate ~init:assocs e) st + | SMatrix (e1, e2) -> Expr.Labelled.(associate ~init:(associate ~init:assocs e1) e2) + | SArray (st, e) -> associate ~init:(Expr.Labelled.associate ~init:assocs e) st +;; -let is_scalar = function SInt | SReal -> true | _ -> false -let rec inner_type = function SArray (t, _) -> inner_type t | t -> t +let is_scalar = function + | SInt | SReal -> true + | _ -> false +;; + +let rec inner_type = function + | SArray (t, _) -> inner_type t + | t -> t +;; let rec dims_of st = match st with | SArray (t, _) -> dims_of t - | SMatrix (d1, d2) -> [d1; d2] - | SRowVector dim | SVector dim -> [dim] + | SMatrix (d1, d2) -> [ d1; d2 ] + | SRowVector dim | SVector dim -> [ dim ] | SInt | SReal -> [] +;; let rec get_dims = function | SInt | SReal -> [] - | SVector d | SRowVector d -> [d] - | SMatrix (dim1, dim2) -> [dim1; dim2] + | SVector d | SRowVector d -> [ d ] + | SMatrix (dim1, dim2) -> [ dim1; dim2 ] | SArray (t, dim) -> dim :: get_dims t +;; let%expect_test "dims" = let open Fmt in - strf "@[%a@]" (list ~sep:comma string) - (get_dims (SArray (SMatrix ("x", "y"), "z"))) - |> print_endline ; + strf "@[%a@]" (list ~sep:comma string) (get_dims (SArray (SMatrix ("x", "y"), "z"))) + |> print_endline; [%expect {| z, x, y |}] +;; diff --git a/src/middle/Stan_math_signatures.ml b/src/middle/Stan_math_signatures.ml index 2b1b488ec6..7b6d80e9ab 100644 --- a/src/middle/Stan_math_signatures.ml +++ b/src/middle/Stan_math_signatures.ml @@ -26,28 +26,39 @@ type dimensionality = just used for element-wise vectorized unary functions now *) let rec bare_array_type (t, i) = - match i with 0 -> t | j -> UnsizedType.UArray (bare_array_type (t, j - 1)) + match i with + | 0 -> t + | j -> UnsizedType.UArray (bare_array_type (t, j - 1)) +;; let rec expand_arg = function - | DReal -> [UnsizedType.UReal] - | DVector -> [UVector] - | DMatrix -> [UMatrix] - | DVInt -> [UInt; UArray UInt] - | DVReal -> [UReal; UArray UReal; UVector; URowVector] + | DReal -> [ UnsizedType.UReal ] + | DVector -> [ UVector ] + | DMatrix -> [ UMatrix ] + | DVInt -> [ UInt; UArray UInt ] + | DVReal -> [ UReal; UArray UReal; UVector; URowVector ] | DIntAndReals -> expand_arg DVReal @ expand_arg DVInt - | DVectors -> [UVector; UArray UVector; URowVector; UArray URowVector] + | DVectors -> [ UVector; UArray UVector; URowVector; UArray URowVector ] | DDeepVectorized -> - let all_base = [UnsizedType.UInt; UReal; URowVector; UVector; UMatrix] in - List.( - concat_map all_base ~f:(fun a -> - map (range 0 8) ~f:(fun i -> bare_array_type (a, i)) )) + let all_base = [ UnsizedType.UInt; UReal; URowVector; UVector; UMatrix ] in + List.( + concat_map all_base ~f:(fun a -> + map (range 0 8) ~f:(fun i -> bare_array_type (a, i)))) +;; -type fkind = Lpmf | Lpdf | Rng | Cdf | Ccdf | UnaryVectorized +type fkind = + | Lpmf + | Lpdf + | Rng + | Cdf + | Ccdf + | UnaryVectorized let is_primitive = function | UnsizedType.UReal -> true | UInt -> true | _ -> false +;; (** The signatures hash table *) let stan_math_signatures = String.Table.create () @@ -58,31 +69,36 @@ let manual_stan_math_signatures = String.Table.create () (* XXX The correct word here isn't combination - what is it? *) let all_combinations xx = - List.fold_right xx ~init:[[]] ~f:(fun x accum -> - List.concat_map accum ~f:(fun acc -> - List.map ~f:(fun arg -> arg :: acc) x ) ) + List.fold_right xx ~init:[ [] ] ~f:(fun x accum -> + List.concat_map accum ~f:(fun acc -> List.map ~f:(fun arg -> arg :: acc) x)) +;; let%expect_test "combinations " = - let a = all_combinations [[1; 2]; [3; 4]; [5; 6]] in - [%sexp (a : int list list)] |> Sexp.to_string_hum |> print_endline ; - [%expect - {| ((1 3 5) (2 3 5) (1 4 5) (2 4 5) (1 3 6) (2 3 6) (1 4 6) (2 4 6)) |}] + let a = all_combinations [ [ 1; 2 ]; [ 3; 4 ]; [ 5; 6 ] ] in + [%sexp (a : int list list)] |> Sexp.to_string_hum |> print_endline; + [%expect {| ((1 3 5) (2 3 5) (1 4 5) (2 4 5) (1 3 6) (2 3 6) (1 4 6) (2 4 6)) |}] +;; -let missing_math_functions = String.Set.of_list ["beta_proportion_cdf"] +let missing_math_functions = String.Set.of_list [ "beta_proportion_cdf" ] let rng_return_type t lt = if List.for_all ~f:is_primitive lt then t else UnsizedType.UArray t +;; let add_unqualified (name, rt, uqargts) = - Hashtbl.add_multi manual_stan_math_signatures ~key:name - ~data:(rt, List.map ~f:(fun x -> (UnsizedType.AutoDiffable, x)) uqargts) + Hashtbl.add_multi + manual_stan_math_signatures + ~key:name + ~data:(rt, List.map ~f:(fun x -> UnsizedType.AutoDiffable, x) uqargts) +;; let rec ints_to_real = function | UnsizedType.UInt -> UnsizedType.UReal | UArray t -> UArray (ints_to_real t) | x -> x +;; -let reduce_sum_allowed_dimensionalities = [1; 2; 3; 4; 5; 6; 7] +let reduce_sum_allowed_dimensionalities = [ 1; 2; 3; 4; 5; 6; 7 ] let reduce_sum_slice_types = let base_slice_type i = @@ -90,20 +106,25 @@ let reduce_sum_slice_types = ; bare_array_type (UnsizedType.UInt, i) ; bare_array_type (UnsizedType.UMatrix, i) ; bare_array_type (UnsizedType.UVector, i) - ; bare_array_type (UnsizedType.URowVector, i) ] + ; bare_array_type (UnsizedType.URowVector, i) + ] in List.concat (List.map ~f:base_slice_type reduce_sum_allowed_dimensionalities) +;; let mk_declarative_sig (fnkinds, name, args) = let sfxes = function - | Lpmf -> ["_lpmf"; "_log"] - | Lpdf -> ["_lpdf"; "_log"] - | Rng -> ["_rng"] - | Cdf -> ["_cdf"; "_cdf_log"; "_lcdf"] - | Ccdf -> ["_ccdf_log"; "_lccdf"] - | UnaryVectorized -> [""] + | Lpmf -> [ "_lpmf"; "_log" ] + | Lpdf -> [ "_lpdf"; "_log" ] + | Rng -> [ "_rng" ] + | Cdf -> [ "_cdf"; "_cdf_log"; "_lcdf" ] + | Ccdf -> [ "_ccdf_log"; "_lccdf" ] + | UnaryVectorized -> [ "" ] + in + let add_ints = function + | DVReal -> DIntAndReals + | x -> x in - let add_ints = function DVReal -> DIntAndReals | x -> x in let all_expanded args = all_combinations (List.map ~f:expand_arg args) in let promoted_dim = function | DVInt -> UnsizedType.UInt @@ -117,136 +138,135 @@ let mk_declarative_sig (fnkinds, name, args) = in let create_from_fk_args fk arglists = List.concat_map arglists ~f:(fun args -> - List.map (sfxes fk) ~f:(fun sfx -> - (name ^ sfx, find_rt UReal args fk, args) ) ) + List.map (sfxes fk) ~f:(fun sfx -> name ^ sfx, find_rt UReal args fk, args)) in let add_fnkind = function | Rng -> - let rt, args = (List.hd_exn args, List.tl_exn args) in - let args = List.map ~f:add_ints args in - let rt = promoted_dim rt in - let name = name ^ "_rng" in - List.map (all_expanded args) ~f:(fun args -> - (name, find_rt rt args Rng, args) ) - | UnaryVectorized -> - create_from_fk_args UnaryVectorized (all_expanded args) + let rt, args = List.hd_exn args, List.tl_exn args in + let args = List.map ~f:add_ints args in + let rt = promoted_dim rt in + let name = name ^ "_rng" in + List.map (all_expanded args) ~f:(fun args -> name, find_rt rt args Rng, args) + | UnaryVectorized -> create_from_fk_args UnaryVectorized (all_expanded args) | fk -> create_from_fk_args fk (all_expanded args) in List.concat_map fnkinds ~f:add_fnkind |> List.filter ~f:(fun (n, _, _) -> not (Set.mem missing_math_functions n)) |> List.map ~f:(fun (n, rt, args) -> - (n, rt, List.map ~f:(fun x -> (UnsizedType.AutoDiffable, x)) args) ) + n, rt, List.map ~f:(fun x -> UnsizedType.AutoDiffable, x) args) +;; -let full_lpdf = [Lpdf; Rng; Ccdf; Cdf] -let full_lpmf = [Lpmf; Rng; Ccdf; Cdf] -let reduce_sum_functions = ["reduce_sum"; "reduce_sum_static"] +let full_lpdf = [ Lpdf; Rng; Ccdf; Cdf ] +let full_lpmf = [ Lpmf; Rng; Ccdf; Cdf ] +let reduce_sum_functions = [ "reduce_sum"; "reduce_sum_static" ] let is_reduce_sum_fn f = List.mem ~equal:String.equal reduce_sum_functions f let distributions = - [ (full_lpmf, "beta_binomial", [DVInt; DVInt; DVReal; DVReal]) - ; (full_lpdf, "beta", [DVReal; DVReal; DVReal]) - ; ([Lpdf; Ccdf; Cdf], "beta_proportion", [DVReal; DVReal; DIntAndReals]) - ; (full_lpmf, "bernoulli", [DVInt; DVReal]) - ; ([Lpmf; Rng], "bernoulli_logit", [DVInt; DVReal]) - ; (full_lpmf, "binomial", [DVInt; DVInt; DVReal]) - ; ([Lpmf], "binomial_logit", [DVInt; DVInt; DVReal]) - ; ([Lpmf], "categorical", [DVInt; DVector]) - ; ([Lpmf], "categorical_logit", [DVInt; DVector]) - ; (full_lpdf, "cauchy", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "chi_square", [DVReal; DVReal]) - ; ([Lpdf], "dirichlet", [DVectors; DVectors]) - ; (full_lpdf, "double_exponential", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "exp_mod_normal", [DVReal; DVReal; DVReal; DVReal]) - ; (full_lpdf, "exponential", [DVReal; DVReal]) - ; (full_lpdf, "frechet", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "gamma", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "gumbel", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "inv_chi_square", [DVReal; DVReal]) - ; (full_lpdf, "inv_gamma", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "logistic", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "lognormal", [DVReal; DVReal; DVReal]) - ; ([Lpdf], "multi_gp", [DMatrix; DMatrix; DVector]) - ; ([Lpdf], "multi_gp_cholesky", [DMatrix; DMatrix; DVector]) - ; ([Lpdf], "multi_normal", [DVectors; DVectors; DMatrix]) - ; ([Lpdf], "multi_normal_cholesky", [DVectors; DVectors; DMatrix]) - ; ([Lpdf], "multi_normal_prec", [DVectors; DVectors; DMatrix]) - ; ([Lpdf], "multi_student_t", [DVectors; DReal; DVectors; DMatrix]) - ; (full_lpmf, "neg_binomial", [DVInt; DVReal; DVReal]) - ; (full_lpmf, "neg_binomial_2", [DVInt; DVReal; DVReal]) - ; ([Lpmf; Rng], "neg_binomial_2_log", [DVInt; DVReal; DVReal]) - ; (full_lpdf, "normal", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "pareto", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "pareto_type_2", [DVReal; DVReal; DVReal; DVReal]) - ; (full_lpmf, "poisson", [DVInt; DVReal]) - ; ([Lpmf; Rng], "poisson_log", [DVInt; DVReal]) - ; (full_lpdf, "rayleigh", [DVReal; DVReal]) - ; (full_lpdf, "scaled_inv_chi_square", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "skew_normal", [DVReal; DVReal; DVReal; DVReal]) - ; (full_lpdf, "student_t", [DVReal; DVReal; DVReal; DVReal]) - ; (full_lpdf, "std_normal", [DVReal]) - ; (full_lpdf, "uniform", [DVReal; DVReal; DVReal]) - ; ([Lpdf; Rng], "von_mises", [DVReal; DVReal; DVReal]) - ; (full_lpdf, "weibull", [DVReal; DVReal; DVReal]) - ; ([Lpdf], "wiener", [DVReal; DVReal; DVReal; DVReal; DVReal]) - ; ([Lpdf], "wishart", [DMatrix; DReal; DMatrix]) ] + [ full_lpmf, "beta_binomial", [ DVInt; DVInt; DVReal; DVReal ] + ; full_lpdf, "beta", [ DVReal; DVReal; DVReal ] + ; [ Lpdf; Ccdf; Cdf ], "beta_proportion", [ DVReal; DVReal; DIntAndReals ] + ; full_lpmf, "bernoulli", [ DVInt; DVReal ] + ; [ Lpmf; Rng ], "bernoulli_logit", [ DVInt; DVReal ] + ; full_lpmf, "binomial", [ DVInt; DVInt; DVReal ] + ; [ Lpmf ], "binomial_logit", [ DVInt; DVInt; DVReal ] + ; [ Lpmf ], "categorical", [ DVInt; DVector ] + ; [ Lpmf ], "categorical_logit", [ DVInt; DVector ] + ; full_lpdf, "cauchy", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "chi_square", [ DVReal; DVReal ] + ; [ Lpdf ], "dirichlet", [ DVectors; DVectors ] + ; full_lpdf, "double_exponential", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "exp_mod_normal", [ DVReal; DVReal; DVReal; DVReal ] + ; full_lpdf, "exponential", [ DVReal; DVReal ] + ; full_lpdf, "frechet", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "gamma", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "gumbel", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "inv_chi_square", [ DVReal; DVReal ] + ; full_lpdf, "inv_gamma", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "logistic", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "lognormal", [ DVReal; DVReal; DVReal ] + ; [ Lpdf ], "multi_gp", [ DMatrix; DMatrix; DVector ] + ; [ Lpdf ], "multi_gp_cholesky", [ DMatrix; DMatrix; DVector ] + ; [ Lpdf ], "multi_normal", [ DVectors; DVectors; DMatrix ] + ; [ Lpdf ], "multi_normal_cholesky", [ DVectors; DVectors; DMatrix ] + ; [ Lpdf ], "multi_normal_prec", [ DVectors; DVectors; DMatrix ] + ; [ Lpdf ], "multi_student_t", [ DVectors; DReal; DVectors; DMatrix ] + ; full_lpmf, "neg_binomial", [ DVInt; DVReal; DVReal ] + ; full_lpmf, "neg_binomial_2", [ DVInt; DVReal; DVReal ] + ; [ Lpmf; Rng ], "neg_binomial_2_log", [ DVInt; DVReal; DVReal ] + ; full_lpdf, "normal", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "pareto", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "pareto_type_2", [ DVReal; DVReal; DVReal; DVReal ] + ; full_lpmf, "poisson", [ DVInt; DVReal ] + ; [ Lpmf; Rng ], "poisson_log", [ DVInt; DVReal ] + ; full_lpdf, "rayleigh", [ DVReal; DVReal ] + ; full_lpdf, "scaled_inv_chi_square", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "skew_normal", [ DVReal; DVReal; DVReal; DVReal ] + ; full_lpdf, "student_t", [ DVReal; DVReal; DVReal; DVReal ] + ; full_lpdf, "std_normal", [ DVReal ] + ; full_lpdf, "uniform", [ DVReal; DVReal; DVReal ] + ; [ Lpdf; Rng ], "von_mises", [ DVReal; DVReal; DVReal ] + ; full_lpdf, "weibull", [ DVReal; DVReal; DVReal ] + ; [ Lpdf ], "wiener", [ DVReal; DVReal; DVReal; DVReal; DVReal ] + ; [ Lpdf ], "wishart", [ DMatrix; DReal; DMatrix ] + ] +;; let math_sigs = - [ ([UnaryVectorized], "acos", [DDeepVectorized]) - ; ([UnaryVectorized], "acosh", [DDeepVectorized]) - ; ([UnaryVectorized], "asin", [DDeepVectorized]) - ; ([UnaryVectorized], "asinh", [DDeepVectorized]) - ; ([UnaryVectorized], "atan", [DDeepVectorized]) - ; ([UnaryVectorized], "atanh", [DDeepVectorized]) - ; ([UnaryVectorized], "cbrt", [DDeepVectorized]) - ; ([UnaryVectorized], "ceil", [DDeepVectorized]) - ; ([UnaryVectorized], "cos", [DDeepVectorized]) - ; ([UnaryVectorized], "cosh", [DDeepVectorized]) - ; ([UnaryVectorized], "digamma", [DDeepVectorized]) - ; ([UnaryVectorized], "erf", [DDeepVectorized]) - ; ([UnaryVectorized], "erfc", [DDeepVectorized]) - ; ([UnaryVectorized], "exp", [DDeepVectorized]) - ; ([UnaryVectorized], "exp2", [DDeepVectorized]) - ; ([UnaryVectorized], "expm1", [DDeepVectorized]) - ; ([UnaryVectorized], "fabs", [DDeepVectorized]) - ; ([UnaryVectorized], "floor", [DDeepVectorized]) - ; ([UnaryVectorized], "inv", [DDeepVectorized]) - ; ([UnaryVectorized], "inv_cloglog", [DDeepVectorized]) - ; ([UnaryVectorized], "inv_logit", [DDeepVectorized]) - ; ([UnaryVectorized], "inv_Phi", [DDeepVectorized]) - ; ([UnaryVectorized], "inv_sqrt", [DDeepVectorized]) - ; ([UnaryVectorized], "inv_square", [DDeepVectorized]) - ; ([UnaryVectorized], "lambert_w0", [DDeepVectorized]) - ; ([UnaryVectorized], "lambert_wm1", [DDeepVectorized]) - ; ([UnaryVectorized], "lgamma", [DDeepVectorized]) - ; ([UnaryVectorized], "log", [DDeepVectorized]) - ; ([UnaryVectorized], "log10", [DDeepVectorized]) - ; ([UnaryVectorized], "log1m", [DDeepVectorized]) - ; ([UnaryVectorized], "log1m_exp", [DDeepVectorized]) - ; ([UnaryVectorized], "log1m_inv_logit", [DDeepVectorized]) - ; ([UnaryVectorized], "log1p", [DDeepVectorized]) - ; ([UnaryVectorized], "log1p_exp", [DDeepVectorized]) - ; ([UnaryVectorized], "log2", [DDeepVectorized]) - ; ([UnaryVectorized], "log_inv_logit", [DDeepVectorized]) - ; ([UnaryVectorized], "logit", [DDeepVectorized]) - ; ([UnaryVectorized], "Phi", [DDeepVectorized]) - ; ([UnaryVectorized], "Phi_approx", [DDeepVectorized]) - ; ([UnaryVectorized], "round", [DDeepVectorized]) - ; ([UnaryVectorized], "sin", [DDeepVectorized]) - ; ([UnaryVectorized], "sinh", [DDeepVectorized]) - ; ([UnaryVectorized], "sqrt", [DDeepVectorized]) - ; ([UnaryVectorized], "square", [DDeepVectorized]) - ; ([UnaryVectorized], "step", [DReal]) - ; ([UnaryVectorized], "tan", [DDeepVectorized]) - ; ([UnaryVectorized], "tanh", [DDeepVectorized]) - (* ; add_nullary ("target") *) - ; ([UnaryVectorized], "tgamma", [DDeepVectorized]) - ; ([UnaryVectorized], "trunc", [DDeepVectorized]) - ; ([UnaryVectorized], "trigamma", [DDeepVectorized]) ] + [ [ UnaryVectorized ], "acos", [ DDeepVectorized ] + ; [ UnaryVectorized ], "acosh", [ DDeepVectorized ] + ; [ UnaryVectorized ], "asin", [ DDeepVectorized ] + ; [ UnaryVectorized ], "asinh", [ DDeepVectorized ] + ; [ UnaryVectorized ], "atan", [ DDeepVectorized ] + ; [ UnaryVectorized ], "atanh", [ DDeepVectorized ] + ; [ UnaryVectorized ], "cbrt", [ DDeepVectorized ] + ; [ UnaryVectorized ], "ceil", [ DDeepVectorized ] + ; [ UnaryVectorized ], "cos", [ DDeepVectorized ] + ; [ UnaryVectorized ], "cosh", [ DDeepVectorized ] + ; [ UnaryVectorized ], "digamma", [ DDeepVectorized ] + ; [ UnaryVectorized ], "erf", [ DDeepVectorized ] + ; [ UnaryVectorized ], "erfc", [ DDeepVectorized ] + ; [ UnaryVectorized ], "exp", [ DDeepVectorized ] + ; [ UnaryVectorized ], "exp2", [ DDeepVectorized ] + ; [ UnaryVectorized ], "expm1", [ DDeepVectorized ] + ; [ UnaryVectorized ], "fabs", [ DDeepVectorized ] + ; [ UnaryVectorized ], "floor", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv_cloglog", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv_logit", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv_Phi", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv_sqrt", [ DDeepVectorized ] + ; [ UnaryVectorized ], "inv_square", [ DDeepVectorized ] + ; [ UnaryVectorized ], "lambert_w0", [ DDeepVectorized ] + ; [ UnaryVectorized ], "lambert_wm1", [ DDeepVectorized ] + ; [ UnaryVectorized ], "lgamma", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log10", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log1m", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log1m_exp", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log1m_inv_logit", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log1p", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log1p_exp", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log2", [ DDeepVectorized ] + ; [ UnaryVectorized ], "log_inv_logit", [ DDeepVectorized ] + ; [ UnaryVectorized ], "logit", [ DDeepVectorized ] + ; [ UnaryVectorized ], "Phi", [ DDeepVectorized ] + ; [ UnaryVectorized ], "Phi_approx", [ DDeepVectorized ] + ; [ UnaryVectorized ], "round", [ DDeepVectorized ] + ; [ UnaryVectorized ], "sin", [ DDeepVectorized ] + ; [ UnaryVectorized ], "sinh", [ DDeepVectorized ] + ; [ UnaryVectorized ], "sqrt", [ DDeepVectorized ] + ; [ UnaryVectorized ], "square", [ DDeepVectorized ] + ; [ UnaryVectorized ], "step", [ DReal ] + ; [ UnaryVectorized ], "tan", [ DDeepVectorized ] + ; [ UnaryVectorized ], "tanh", [ DDeepVectorized ] (* ; add_nullary ("target") *) + ; [ UnaryVectorized ], "tgamma", [ DDeepVectorized ] + ; [ UnaryVectorized ], "trunc", [ DDeepVectorized ] + ; [ UnaryVectorized ], "trigamma", [ DDeepVectorized ] + ] +;; let all_declarative_sigs = distributions @ math_sigs - -let declarative_fnsigs = - List.concat_map ~f:mk_declarative_sig all_declarative_sigs +let declarative_fnsigs = List.concat_map ~f:mk_declarative_sig all_declarative_sigs (* -- Querying stan_math_signatures -- *) let stan_math_returntype name args = @@ -254,24 +274,28 @@ let stan_math_returntype name args = let namematches = Hashtbl.find_multi stan_math_signatures name in let filteredmatches = List.filter - ~f:(fun x -> - UnsizedType.check_compatible_arguments_mod_conv name (snd x) args ) + ~f:(fun x -> UnsizedType.check_compatible_arguments_mod_conv name (snd x) args) namematches in match name with | x when is_reduce_sum_fn x -> Some (UnsizedType.ReturnType UReal) | _ -> - if List.length filteredmatches = 0 then None - (* Return the least return type in case there are multiple options (due to implicit UInt-UReal conversion), where UInt Some "assign_add" @@ -281,81 +305,93 @@ let assignmentoperator_to_stan_math_fn = function | EltTimes -> Some "assign_elt_times" | EltDivide -> Some "assign_elt_divide" | _ -> None +;; let assignmentoperator_stan_math_return_type assop arg_tys = assignmentoperator_to_stan_math_fn assop |> Option.bind ~f:(fun name -> stan_math_returntype name arg_tys) +;; let operator_to_stan_math_fns = function - | Operator.Plus -> ["add"] - | PPlus -> ["plus"] - | Minus -> ["subtract"] - | PMinus -> ["minus"] - | Times -> ["multiply"] - | Divide -> ["mdivide_right"; "divide"] - | Modulo -> ["modulus"] + | Operator.Plus -> [ "add" ] + | PPlus -> [ "plus" ] + | Minus -> [ "subtract" ] + | PMinus -> [ "minus" ] + | Times -> [ "multiply" ] + | Divide -> [ "mdivide_right"; "divide" ] + | Modulo -> [ "modulus" ] | IntDivide -> [] - | LDivide -> ["mdivide_left"] - | EltTimes -> ["elt_multiply"] - | EltDivide -> ["elt_divide"] - | Pow -> ["pow"] - | Or -> ["logical_or"] - | And -> ["logical_and"] - | Equals -> ["logical_eq"] - | NEquals -> ["logical_neq"] - | Less -> ["logical_lt"] - | Leq -> ["logical_lte"] - | Greater -> ["logical_gt"] - | Geq -> ["logical_gte"] - | PNot -> ["logical_negation"] - | Transpose -> ["transpose"] + | LDivide -> [ "mdivide_left" ] + | EltTimes -> [ "elt_multiply" ] + | EltDivide -> [ "elt_divide" ] + | Pow -> [ "pow" ] + | Or -> [ "logical_or" ] + | And -> [ "logical_and" ] + | Equals -> [ "logical_eq" ] + | NEquals -> [ "logical_neq" ] + | Less -> [ "logical_lt" ] + | Leq -> [ "logical_lte" ] + | Greater -> [ "logical_gt" ] + | Geq -> [ "logical_gte" ] + | PNot -> [ "logical_negation" ] + | Transpose -> [ "transpose" ] +;; let int_divide_type = - UnsizedType.(ReturnType UInt, [(AutoDiffable, UInt); (AutoDiffable, UInt)]) + UnsizedType.(ReturnType UInt, [ AutoDiffable, UInt; AutoDiffable, UInt ]) +;; let operator_stan_math_return_type op arg_tys = - match (op, arg_tys) with - | Operator.IntDivide, [(_, UnsizedType.UInt); (_, UInt)] -> - Some UnsizedType.(ReturnType UInt) + match op, arg_tys with + | Operator.IntDivide, [ (_, UnsizedType.UInt); (_, UInt) ] -> + Some UnsizedType.(ReturnType UInt) | IntDivide, _ -> None | _ -> - operator_to_stan_math_fns op - |> List.filter_map ~f:(fun name -> stan_math_returntype name arg_tys) - |> List.hd + operator_to_stan_math_fns op + |> List.filter_map ~f:(fun name -> stan_math_returntype name arg_tys) + |> List.hd +;; let get_sigs name = let name = Utils.stdlib_distribution_name name in Hashtbl.find_multi stan_math_signatures name |> List.sort ~compare +;; let pp_math_sig ppf (rt, args) = UnsizedType.pp ppf (UFun (args, rt)) - -let pp_math_sigs ppf name = - (Fmt.list ~sep:Fmt.cut pp_math_sig) ppf (get_sigs name) - +let pp_math_sigs ppf name = (Fmt.list ~sep:Fmt.cut pp_math_sig) ppf (get_sigs name) let pretty_print_math_sigs = Fmt.strf "@[@,%a@]" pp_math_sigs let pretty_print_all_math_sigs ppf () = let open Fmt in let pp_sig ppf (name, (rt, args)) = - pf ppf "%a %s(@[%a@])" UnsizedType.pp_returntype rt name + pf + ppf + "%a %s(@[%a@])" + UnsizedType.pp_returntype + rt + name (list ~sep:comma UnsizedType.pp) (List.map ~f:snd args) in let pp_sigs_for_name ppf name = - (list ~sep:cut pp_sig) ppf - (List.map ~f:(fun t -> (name, t)) (get_sigs name)) + (list ~sep:cut pp_sig) ppf (List.map ~f:(fun t -> name, t) (get_sigs name)) in - pf ppf "@[%a@]" + pf + ppf + "@[%a@]" (list ~sep:cut pp_sigs_for_name) (List.sort ~compare (Hashtbl.keys stan_math_signatures)) +;; let pretty_print_math_lib_operator_sigs op = - if op = Operator.IntDivide then - [Fmt.strf "@[@,%a@]" pp_math_sig int_divide_type] + if op = Operator.IntDivide + then [ Fmt.strf "@[@,%a@]" pp_math_sig int_divide_type ] else operator_to_stan_math_fns op |> List.map ~f:pretty_print_math_sigs +;; let pretty_print_math_lib_assignmentoperator_sigs op = assignmentoperator_to_stan_math_fn op |> Option.map ~f:pretty_print_math_sigs +;; (* -- Some helper definitions to populate stan_math_signatures -- *) let bare_types = function @@ -365,6 +401,7 @@ let bare_types = function | 3 -> URowVector | 4 -> UMatrix | i -> raise_s [%sexp (i : int)] +;; let bare_types_size = 5 @@ -374,6 +411,7 @@ let vector_types = function | 2 -> UVector | 3 -> URowVector | i -> raise_s [%sexp (i : int)] +;; let vector_types_size = 4 @@ -381,6 +419,7 @@ let primitive_types = function | 0 -> UnsizedType.UInt | 1 -> UReal | i -> raise_s [%sexp (i : int)] +;; let primitive_types_size = 2 @@ -392,1387 +431,1338 @@ let all_vector_types = function | 4 -> UInt | 5 -> UArray UInt | i -> raise_s [%sexp (i : int)] +;; let all_vector_types_size = 6 let add_qualified (name, rt, argts) = Hashtbl.add_multi stan_math_signatures ~key:name ~data:(rt, argts) +;; let add_nullary name = add_unqualified (name, UnsizedType.ReturnType UReal, []) let add_binary name = - add_unqualified (name, ReturnType UReal, [UnsizedType.UReal; UReal]) + add_unqualified (name, ReturnType UReal, [ UnsizedType.UReal; UReal ]) +;; -let add_ternary name = - add_unqualified (name, ReturnType UReal, [UReal; UReal; UReal]) +let add_ternary name = add_unqualified (name, ReturnType UReal, [ UReal; UReal; UReal ]) let for_all_vector_types s = for i = 0 to all_vector_types_size - 1 do s (all_vector_types i) done +;; let for_vector_types s = for i = 0 to vector_types_size - 1 do s (vector_types i) done +;; (* -- Start populating stan_math_signaturess -- *) let () = List.iter declarative_fnsigs ~f:(fun (key, rt, args) -> - Hashtbl.add_multi stan_math_signatures ~key ~data:(rt, args) ) ; - add_unqualified ("abs", ReturnType UInt, [UInt]) ; - add_unqualified ("abs", ReturnType UReal, [UReal]) ; + Hashtbl.add_multi stan_math_signatures ~key ~data:(rt, args)); + add_unqualified ("abs", ReturnType UInt, [ UInt ]); + add_unqualified ("abs", ReturnType UReal, [ UReal ]); for i = 0 to bare_types_size - 1 do - add_unqualified - ("add", ReturnType (bare_types i), [bare_types i; bare_types i]) - done ; - add_unqualified ("add", ReturnType UVector, [UVector; UReal]) ; - add_unqualified ("add", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("add", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("add", ReturnType UVector, [UReal; UVector]) ; - add_unqualified ("add", ReturnType URowVector, [UReal; URowVector]) ; - add_unqualified ("add", ReturnType UMatrix, [UReal; UMatrix]) ; - add_unqualified ("add_diag", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("add_diag", ReturnType UMatrix, [UMatrix; UVector]) ; - add_unqualified ("add_diag", ReturnType UMatrix, [UMatrix; URowVector]) ; + add_unqualified ("add", ReturnType (bare_types i), [ bare_types i; bare_types i ]) + done; + add_unqualified ("add", ReturnType UVector, [ UVector; UReal ]); + add_unqualified ("add", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("add", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("add", ReturnType UVector, [ UReal; UVector ]); + add_unqualified ("add", ReturnType URowVector, [ UReal; URowVector ]); + add_unqualified ("add", ReturnType UMatrix, [ UReal; UMatrix ]); + add_unqualified ("add_diag", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("add_diag", ReturnType UMatrix, [ UMatrix; UVector ]); + add_unqualified ("add_diag", ReturnType UMatrix, [ UMatrix; URowVector ]); add_qualified ( "algebra_solver" , ReturnType UVector , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UVector ) ) - ; (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "algebra_solver" , ReturnType UVector , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UVector ) ) - ; (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) - ; (DataOnly, UReal); (DataOnly, UReal) ] ) ; + ; AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ; DataOnly, UReal + ; DataOnly, UReal + ] ); add_qualified ( "algebra_solver_newton" , ReturnType UVector , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UVector ) ) - ; (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "algebra_solver_newton" , ReturnType UVector , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UVector ) ) - ; (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) - ; (DataOnly, UReal); (DataOnly, UReal) ] ) ; + ; AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ; DataOnly, UReal + ; DataOnly, UReal + ] ); for i = 1 to 8 - 1 do add_unqualified ( "append_array" , ReturnType (bare_array_type (UInt, i)) - , [bare_array_type (UInt, i); bare_array_type (UInt, i)] ) ; + , [ bare_array_type (UInt, i); bare_array_type (UInt, i) ] ); add_unqualified ( "append_array" , ReturnType (bare_array_type (UReal, i)) - , [bare_array_type (UReal, i); bare_array_type (UReal, i)] ) ; + , [ bare_array_type (UReal, i); bare_array_type (UReal, i) ] ); add_unqualified ( "append_array" , ReturnType (bare_array_type (UVector, i)) - , [bare_array_type (UVector, i); bare_array_type (UVector, i)] ) ; + , [ bare_array_type (UVector, i); bare_array_type (UVector, i) ] ); add_unqualified ( "append_array" , ReturnType (bare_array_type (URowVector, i)) - , [bare_array_type (URowVector, i); bare_array_type (URowVector, i)] ) ; + , [ bare_array_type (URowVector, i); bare_array_type (URowVector, i) ] ); add_unqualified ( "append_array" , ReturnType (bare_array_type (UMatrix, i)) - , [bare_array_type (UMatrix, i); bare_array_type (UMatrix, i)] ) - done ; - add_unqualified ("assign_multiply", Void, [UInt; UInt]) ; - add_unqualified ("assign_multiply", Void, [UMatrix; UMatrix]) ; - add_unqualified ("assign_multiply", Void, [UMatrix; UReal]) ; - add_unqualified ("assign_multiply", Void, [UReal; UReal]) ; - add_unqualified ("assign_multiply", Void, [URowVector; UReal]) ; - add_unqualified ("assign_multiply", Void, [UMatrix; UInt]) ; - add_unqualified ("assign_multiply", Void, [UReal; UInt]) ; - add_unqualified ("assign_multiply", Void, [URowVector; UInt]) ; - add_unqualified ("assign_multiply", Void, [URowVector; UMatrix]) ; - add_unqualified ("assign_multiply", Void, [UVector; UReal]) ; - add_unqualified ("assign_multiply", Void, [UVector; UInt]) ; - add_unqualified ("assign_add", Void, [UInt; UInt]) ; - add_unqualified ("assign_add", Void, [UMatrix; UMatrix]) ; - add_unqualified ("assign_add", Void, [UMatrix; UReal]) ; - add_unqualified ("assign_add", Void, [UReal; UReal]) ; - add_unqualified ("assign_add", Void, [URowVector; UReal]) ; - add_unqualified ("assign_add", Void, [UMatrix; UInt]) ; - add_unqualified ("assign_add", Void, [UReal; UInt]) ; - add_unqualified ("assign_add", Void, [URowVector; UInt]) ; - add_unqualified ("assign_add", Void, [URowVector; URowVector]) ; - add_unqualified ("assign_add", Void, [UVector; UReal]) ; - add_unqualified ("assign_add", Void, [UVector; UInt]) ; - add_unqualified ("assign_add", Void, [UVector; UVector]) ; - add_unqualified ("assign_subtract", Void, [UInt; UInt]) ; - add_unqualified ("assign_subtract", Void, [UMatrix; UMatrix]) ; - add_unqualified ("assign_subtract", Void, [UMatrix; UReal]) ; - add_unqualified ("assign_subtract", Void, [UReal; UReal]) ; - add_unqualified ("assign_subtract", Void, [URowVector; UReal]) ; - add_unqualified ("assign_subtract", Void, [UMatrix; UInt]) ; - add_unqualified ("assign_subtract", Void, [UReal; UInt]) ; - add_unqualified ("assign_subtract", Void, [URowVector; UInt]) ; - add_unqualified ("assign_subtract", Void, [URowVector; URowVector]) ; - add_unqualified ("assign_subtract", Void, [UVector; UReal]) ; - add_unqualified ("assign_subtract", Void, [UVector; UInt]) ; - add_unqualified ("assign_subtract", Void, [UVector; UVector]) ; - add_unqualified ("assign_elt_times", Void, [UMatrix; UMatrix]) ; - add_unqualified ("assign_elt_times", Void, [URowVector; URowVector]) ; - add_unqualified ("assign_elt_times", Void, [UVector; UVector]) ; - add_unqualified ("assign_elt_divide", Void, [UMatrix; UMatrix]) ; - add_unqualified ("assign_elt_divide", Void, [UMatrix; UReal]) ; - add_unqualified ("assign_elt_divide", Void, [URowVector; UReal]) ; - add_unqualified ("assign_elt_divide", Void, [UMatrix; UInt]) ; - add_unqualified ("assign_elt_divide", Void, [URowVector; UInt]) ; - add_unqualified ("assign_elt_divide", Void, [URowVector; URowVector]) ; - add_unqualified ("assign_elt_divide", Void, [UVector; UReal]) ; - add_unqualified ("assign_elt_divide", Void, [UVector; UInt]) ; - add_unqualified ("assign_elt_divide", Void, [UVector; UVector]) ; - add_unqualified ("assign_divide", Void, [UInt; UInt]) ; - add_unqualified ("assign_divide", Void, [UMatrix; UReal]) ; - add_unqualified ("assign_divide", Void, [UReal; UReal]) ; - add_unqualified ("assign_divide", Void, [URowVector; UReal]) ; - add_unqualified ("assign_divide", Void, [UVector; UReal]) ; - add_unqualified ("assign_divide", Void, [UMatrix; UInt]) ; - add_unqualified ("assign_divide", Void, [UReal; UInt]) ; - add_unqualified ("assign_divide", Void, [URowVector; UInt]) ; - add_unqualified ("assign_divide", Void, [UVector; UInt]) ; - add_binary "atan2" ; + , [ bare_array_type (UMatrix, i); bare_array_type (UMatrix, i) ] ) + done; + add_unqualified ("assign_multiply", Void, [ UInt; UInt ]); + add_unqualified ("assign_multiply", Void, [ UMatrix; UMatrix ]); + add_unqualified ("assign_multiply", Void, [ UMatrix; UReal ]); + add_unqualified ("assign_multiply", Void, [ UReal; UReal ]); + add_unqualified ("assign_multiply", Void, [ URowVector; UReal ]); + add_unqualified ("assign_multiply", Void, [ UMatrix; UInt ]); + add_unqualified ("assign_multiply", Void, [ UReal; UInt ]); + add_unqualified ("assign_multiply", Void, [ URowVector; UInt ]); + add_unqualified ("assign_multiply", Void, [ URowVector; UMatrix ]); + add_unqualified ("assign_multiply", Void, [ UVector; UReal ]); + add_unqualified ("assign_multiply", Void, [ UVector; UInt ]); + add_unqualified ("assign_add", Void, [ UInt; UInt ]); + add_unqualified ("assign_add", Void, [ UMatrix; UMatrix ]); + add_unqualified ("assign_add", Void, [ UMatrix; UReal ]); + add_unqualified ("assign_add", Void, [ UReal; UReal ]); + add_unqualified ("assign_add", Void, [ URowVector; UReal ]); + add_unqualified ("assign_add", Void, [ UMatrix; UInt ]); + add_unqualified ("assign_add", Void, [ UReal; UInt ]); + add_unqualified ("assign_add", Void, [ URowVector; UInt ]); + add_unqualified ("assign_add", Void, [ URowVector; URowVector ]); + add_unqualified ("assign_add", Void, [ UVector; UReal ]); + add_unqualified ("assign_add", Void, [ UVector; UInt ]); + add_unqualified ("assign_add", Void, [ UVector; UVector ]); + add_unqualified ("assign_subtract", Void, [ UInt; UInt ]); + add_unqualified ("assign_subtract", Void, [ UMatrix; UMatrix ]); + add_unqualified ("assign_subtract", Void, [ UMatrix; UReal ]); + add_unqualified ("assign_subtract", Void, [ UReal; UReal ]); + add_unqualified ("assign_subtract", Void, [ URowVector; UReal ]); + add_unqualified ("assign_subtract", Void, [ UMatrix; UInt ]); + add_unqualified ("assign_subtract", Void, [ UReal; UInt ]); + add_unqualified ("assign_subtract", Void, [ URowVector; UInt ]); + add_unqualified ("assign_subtract", Void, [ URowVector; URowVector ]); + add_unqualified ("assign_subtract", Void, [ UVector; UReal ]); + add_unqualified ("assign_subtract", Void, [ UVector; UInt ]); + add_unqualified ("assign_subtract", Void, [ UVector; UVector ]); + add_unqualified ("assign_elt_times", Void, [ UMatrix; UMatrix ]); + add_unqualified ("assign_elt_times", Void, [ URowVector; URowVector ]); + add_unqualified ("assign_elt_times", Void, [ UVector; UVector ]); + add_unqualified ("assign_elt_divide", Void, [ UMatrix; UMatrix ]); + add_unqualified ("assign_elt_divide", Void, [ UMatrix; UReal ]); + add_unqualified ("assign_elt_divide", Void, [ URowVector; UReal ]); + add_unqualified ("assign_elt_divide", Void, [ UMatrix; UInt ]); + add_unqualified ("assign_elt_divide", Void, [ URowVector; UInt ]); + add_unqualified ("assign_elt_divide", Void, [ URowVector; URowVector ]); + add_unqualified ("assign_elt_divide", Void, [ UVector; UReal ]); + add_unqualified ("assign_elt_divide", Void, [ UVector; UInt ]); + add_unqualified ("assign_elt_divide", Void, [ UVector; UVector ]); + add_unqualified ("assign_divide", Void, [ UInt; UInt ]); + add_unqualified ("assign_divide", Void, [ UMatrix; UReal ]); + add_unqualified ("assign_divide", Void, [ UReal; UReal ]); + add_unqualified ("assign_divide", Void, [ URowVector; UReal ]); + add_unqualified ("assign_divide", Void, [ UVector; UReal ]); + add_unqualified ("assign_divide", Void, [ UMatrix; UInt ]); + add_unqualified ("assign_divide", Void, [ UReal; UInt ]); + add_unqualified ("assign_divide", Void, [ URowVector; UInt ]); + add_unqualified ("assign_divide", Void, [ UVector; UInt ]); + add_binary "atan2"; add_unqualified ( "bernoulli_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UReal; UVector] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UReal; UVector ] ); add_unqualified ( "bernoulli_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UVector; UVector ] ); add_unqualified - ( "bernoulli_logit_glm_lpmf" - , ReturnType UReal - , [UInt; UMatrix; UReal; UVector] ) ; + ("bernoulli_logit_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UReal; UVector ]); add_unqualified - ( "bernoulli_logit_glm_lpmf" - , ReturnType UReal - , [UInt; UMatrix; UVector; UVector] ) ; + ("bernoulli_logit_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UVector; UVector ]); add_unqualified ( "bernoulli_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UReal; UVector] ) ; + , [ bare_array_type (UInt, 1); URowVector; UReal; UVector ] ); add_unqualified ( "bernoulli_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UVector; UVector] ) ; - add_unqualified ("bessel_first_kind", ReturnType UReal, [UInt; UReal]) ; - add_unqualified ("bessel_second_kind", ReturnType UReal, [UInt; UReal]) ; + , [ bare_array_type (UInt, 1); URowVector; UVector; UVector ] ); + add_unqualified ("bessel_first_kind", ReturnType UReal, [ UInt; UReal ]); + add_unqualified ("bessel_second_kind", ReturnType UReal, [ UInt; UReal ]); (* XXX For some reason beta_proportion_rng doesn't take ints as first arg *) for_vector_types (fun t -> for_all_vector_types (fun u -> add_unqualified - ( "beta_proportion_rng" - , ReturnType (rng_return_type UReal [t; u]) - , [t; u] ) ) ) ; - add_unqualified ("binary_log_loss", ReturnType UReal, [UInt; UReal]) ; - add_binary "binomial_coefficient_log" ; - add_unqualified - ("block", ReturnType UMatrix, [UMatrix; UInt; UInt; UInt; UInt]) ; - add_unqualified ("categorical_rng", ReturnType UInt, [UVector]) ; - add_unqualified ("categorical_logit_rng", ReturnType UInt, [UVector]) ; + ("beta_proportion_rng", ReturnType (rng_return_type UReal [ t; u ]), [ t; u ]))); + add_unqualified ("binary_log_loss", ReturnType UReal, [ UInt; UReal ]); + add_binary "binomial_coefficient_log"; + add_unqualified ("block", ReturnType UMatrix, [ UMatrix; UInt; UInt; UInt; UInt ]); + add_unqualified ("categorical_rng", ReturnType UInt, [ UVector ]); + add_unqualified ("categorical_logit_rng", ReturnType UInt, [ UVector ]); add_unqualified ( "categorical_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UVector; UMatrix] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UVector; UMatrix ] ); add_unqualified - ( "categorical_logit_glm_lpmf" - , ReturnType UReal - , [UInt; UMatrix; UVector; UMatrix] ) ; + ("categorical_logit_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UVector; UMatrix ]); add_unqualified ( "categorical_logit_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UVector; UMatrix] ) ; + , [ bare_array_type (UInt, 1); URowVector; UVector; UMatrix ] ); add_unqualified ( "categorical_logit_glm_lpmf" , ReturnType UReal - , [UInt; URowVector; UVector; UMatrix] ) ; - add_unqualified ("append_col", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("append_col", ReturnType UMatrix, [UVector; UMatrix]) ; - add_unqualified ("append_col", ReturnType UMatrix, [UMatrix; UVector]) ; - add_unqualified ("append_col", ReturnType UMatrix, [UVector; UVector]) ; - add_unqualified - ("append_col", ReturnType URowVector, [URowVector; URowVector]) ; - add_unqualified ("append_col", ReturnType URowVector, [UReal; URowVector]) ; - add_unqualified ("append_col", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("cholesky_decompose", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("choose", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("col", ReturnType UVector, [UMatrix; UInt]) ; - add_unqualified ("cols", ReturnType UInt, [UVector]) ; - add_unqualified ("cols", ReturnType UInt, [URowVector]) ; - add_unqualified ("cols", ReturnType UInt, [UMatrix]) ; - add_unqualified - ("columns_dot_product", ReturnType URowVector, [UVector; UVector]) ; - add_unqualified - ("columns_dot_product", ReturnType URowVector, [URowVector; URowVector]) ; - add_unqualified - ("columns_dot_product", ReturnType URowVector, [UMatrix; UMatrix]) ; - add_unqualified ("columns_dot_self", ReturnType URowVector, [UVector]) ; - add_unqualified ("columns_dot_self", ReturnType URowVector, [URowVector]) ; - add_unqualified ("columns_dot_self", ReturnType URowVector, [UMatrix]) ; + , [ UInt; URowVector; UVector; UMatrix ] ); + add_unqualified ("append_col", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("append_col", ReturnType UMatrix, [ UVector; UMatrix ]); + add_unqualified ("append_col", ReturnType UMatrix, [ UMatrix; UVector ]); + add_unqualified ("append_col", ReturnType UMatrix, [ UVector; UVector ]); + add_unqualified ("append_col", ReturnType URowVector, [ URowVector; URowVector ]); + add_unqualified ("append_col", ReturnType URowVector, [ UReal; URowVector ]); + add_unqualified ("append_col", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("cholesky_decompose", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("choose", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("col", ReturnType UVector, [ UMatrix; UInt ]); + add_unqualified ("cols", ReturnType UInt, [ UVector ]); + add_unqualified ("cols", ReturnType UInt, [ URowVector ]); + add_unqualified ("cols", ReturnType UInt, [ UMatrix ]); + add_unqualified ("columns_dot_product", ReturnType URowVector, [ UVector; UVector ]); + add_unqualified + ("columns_dot_product", ReturnType URowVector, [ URowVector; URowVector ]); + add_unqualified ("columns_dot_product", ReturnType URowVector, [ UMatrix; UMatrix ]); + add_unqualified ("columns_dot_self", ReturnType URowVector, [ UVector ]); + add_unqualified ("columns_dot_self", ReturnType URowVector, [ URowVector ]); + add_unqualified ("columns_dot_self", ReturnType URowVector, [ UMatrix ]); + add_unqualified + ("cov_exp_quad", ReturnType UMatrix, [ bare_array_type (UReal, 1); UReal; UReal ]); + add_unqualified + ("cov_exp_quad", ReturnType UMatrix, [ bare_array_type (UVector, 1); UReal; UReal ]); + add_unqualified + ("cov_exp_quad", ReturnType UMatrix, [ bare_array_type (URowVector, 1); UReal; UReal ]); add_unqualified ( "cov_exp_quad" , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified ( "cov_exp_quad" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal] ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "cov_exp_quad" , ReturnType UMatrix - , [bare_array_type (URowVector, 1); UReal; UReal] ) ; - add_unqualified - ( "cov_exp_quad" - , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal] ) ; - add_unqualified - ( "cov_exp_quad" - , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal] - ) ; - add_unqualified - ( "cov_exp_quad" - , ReturnType UMatrix - , [ bare_array_type (URowVector, 1) - ; bare_array_type (URowVector, 1) - ; UReal; UReal ] ) ; - add_unqualified ("crossprod", ReturnType UMatrix, [UMatrix]) ; + , [ bare_array_type (URowVector, 1); bare_array_type (URowVector, 1); UReal; UReal ] + ); + add_unqualified ("crossprod", ReturnType UMatrix, [ UMatrix ]); add_unqualified ( "csr_matrix_times_vector" , ReturnType UVector - , [ UInt; UInt; UVector + , [ UInt + ; UInt + ; UVector ; bare_array_type (UInt, 1) ; bare_array_type (UInt, 1) - ; UVector ] ) ; + ; UVector + ] ); add_unqualified ( "csr_to_dense_matrix" , ReturnType UMatrix - , [ UInt; UInt; UVector - ; bare_array_type (UInt, 1) - ; bare_array_type (UInt, 1) ] ) ; - add_unqualified ("csr_extract_w", ReturnType UVector, [UMatrix]) ; - add_unqualified - ("csr_extract_v", ReturnType (bare_array_type (UInt, 1)), [UMatrix]) ; - add_unqualified - ("csr_extract_u", ReturnType (bare_array_type (UInt, 1)), [UMatrix]) ; + , [ UInt; UInt; UVector; bare_array_type (UInt, 1); bare_array_type (UInt, 1) ] ); + add_unqualified ("csr_extract_w", ReturnType UVector, [ UMatrix ]); + add_unqualified ("csr_extract_v", ReturnType (bare_array_type (UInt, 1)), [ UMatrix ]); + add_unqualified ("csr_extract_u", ReturnType (bare_array_type (UInt, 1)), [ UMatrix ]); add_unqualified ( "cumulative_sum" , ReturnType (bare_array_type (UReal, 1)) - , [bare_array_type (UReal, 1)] ) ; - add_unqualified ("cumulative_sum", ReturnType UVector, [UVector]) ; - add_unqualified ("cumulative_sum", ReturnType URowVector, [URowVector]) ; - add_unqualified ("determinant", ReturnType UReal, [UMatrix]) ; - add_unqualified ("diag_matrix", ReturnType UMatrix, [UVector]) ; - add_unqualified ("diag_post_multiply", ReturnType UMatrix, [UMatrix; UVector]) ; - add_unqualified - ("diag_post_multiply", ReturnType UMatrix, [UMatrix; URowVector]) ; - add_unqualified ("diag_pre_multiply", ReturnType UMatrix, [UVector; UMatrix]) ; - add_unqualified - ("diag_pre_multiply", ReturnType UMatrix, [URowVector; UMatrix]) ; - add_unqualified ("diagonal", ReturnType UVector, [UMatrix]) ; - add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [UInt]) ; - add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [UReal]) ; - add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [UVector]) ; - add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [URowVector]) ; - add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [UMatrix]) ; + , [ bare_array_type (UReal, 1) ] ); + add_unqualified ("cumulative_sum", ReturnType UVector, [ UVector ]); + add_unqualified ("cumulative_sum", ReturnType URowVector, [ URowVector ]); + add_unqualified ("determinant", ReturnType UReal, [ UMatrix ]); + add_unqualified ("diag_matrix", ReturnType UMatrix, [ UVector ]); + add_unqualified ("diag_post_multiply", ReturnType UMatrix, [ UMatrix; UVector ]); + add_unqualified ("diag_post_multiply", ReturnType UMatrix, [ UMatrix; URowVector ]); + add_unqualified ("diag_pre_multiply", ReturnType UMatrix, [ UVector; UMatrix ]); + add_unqualified ("diag_pre_multiply", ReturnType UMatrix, [ URowVector; UMatrix ]); + add_unqualified ("diagonal", ReturnType UVector, [ UMatrix ]); + add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [ UInt ]); + add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [ UReal ]); + add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [ UVector ]); + add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [ URowVector ]); + add_unqualified ("dims", ReturnType (bare_array_type (UInt, 1)), [ UMatrix ]); for i = 0 to 8 - 1 do add_unqualified - ( "dims" - , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, i + 1)] ) ; + ("dims", ReturnType (bare_array_type (UInt, 1)), [ bare_array_type (UInt, i + 1) ]); add_unqualified - ( "dims" - , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UReal, i + 1)] ) ; + ("dims", ReturnType (bare_array_type (UInt, 1)), [ bare_array_type (UReal, i + 1) ]); add_unqualified ( "dims" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UVector, i + 1)] ) ; + , [ bare_array_type (UVector, i + 1) ] ); add_unqualified ( "dims" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (URowVector, i + 1)] ) ; + , [ bare_array_type (URowVector, i + 1) ] ); add_unqualified ( "dims" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UMatrix, i + 1)] ) - done ; - add_unqualified ("dirichlet_rng", ReturnType UVector, [UVector]) ; - add_unqualified ("distance", ReturnType UReal, [UVector; UVector]) ; - add_unqualified ("distance", ReturnType UReal, [URowVector; URowVector]) ; - add_unqualified ("distance", ReturnType UReal, [UVector; URowVector]) ; - add_unqualified ("distance", ReturnType UReal, [URowVector; UVector]) ; - add_unqualified ("divide", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("divide", ReturnType UReal, [UReal; UReal]) ; - add_unqualified ("divide", ReturnType UVector, [UVector; UReal]) ; - add_unqualified ("divide", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("divide", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("dot_product", ReturnType UReal, [UVector; UVector]) ; - add_unqualified ("dot_product", ReturnType UReal, [URowVector; URowVector]) ; - add_unqualified ("dot_product", ReturnType UReal, [UVector; URowVector]) ; - add_unqualified ("dot_product", ReturnType UReal, [URowVector; UVector]) ; + , [ bare_array_type (UMatrix, i + 1) ] ) + done; + add_unqualified ("dirichlet_rng", ReturnType UVector, [ UVector ]); + add_unqualified ("distance", ReturnType UReal, [ UVector; UVector ]); + add_unqualified ("distance", ReturnType UReal, [ URowVector; URowVector ]); + add_unqualified ("distance", ReturnType UReal, [ UVector; URowVector ]); + add_unqualified ("distance", ReturnType UReal, [ URowVector; UVector ]); + add_unqualified ("divide", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("divide", ReturnType UReal, [ UReal; UReal ]); + add_unqualified ("divide", ReturnType UVector, [ UVector; UReal ]); + add_unqualified ("divide", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("divide", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("dot_product", ReturnType UReal, [ UVector; UVector ]); + add_unqualified ("dot_product", ReturnType UReal, [ URowVector; URowVector ]); + add_unqualified ("dot_product", ReturnType UReal, [ UVector; URowVector ]); + add_unqualified ("dot_product", ReturnType UReal, [ URowVector; UVector ]); add_unqualified ( "dot_product" , ReturnType UReal - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1)] ) ; - add_unqualified ("dot_self", ReturnType UReal, [UVector]) ; - add_unqualified ("dot_self", ReturnType UReal, [URowVector]) ; - add_nullary "e" ; - add_unqualified ("eigenvalues_sym", ReturnType UVector, [UMatrix]) ; - add_unqualified ("eigenvectors_sym", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("qr_Q", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("qr_R", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("qr_thin_Q", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("qr_thin_R", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("elt_divide", ReturnType UVector, [UVector; UVector]) ; - add_unqualified - ("elt_divide", ReturnType URowVector, [URowVector; URowVector]) ; - add_unqualified ("elt_divide", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("elt_divide", ReturnType UVector, [UVector; UReal]) ; - add_unqualified ("elt_divide", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("elt_divide", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("elt_divide", ReturnType UVector, [UReal; UVector]) ; - add_unqualified ("elt_divide", ReturnType URowVector, [UReal; URowVector]) ; - add_unqualified ("elt_divide", ReturnType UMatrix, [UReal; UMatrix]) ; - add_unqualified ("elt_multiply", ReturnType UVector, [UVector; UVector]) ; - add_unqualified - ("elt_multiply", ReturnType URowVector, [URowVector; URowVector]) ; - add_unqualified ("elt_multiply", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("falling_factorial", ReturnType UReal, [UReal; UInt]) ; - add_unqualified ("falling_factorial", ReturnType UInt, [UInt; UInt]) ; - add_binary "fdim" ; - add_ternary "fma" ; - add_binary "fmax" ; - add_binary "fmin" ; - add_binary "fmod" ; - add_binary "gamma_p" ; - add_binary "gamma_q" ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1) ] ); + add_unqualified ("dot_self", ReturnType UReal, [ UVector ]); + add_unqualified ("dot_self", ReturnType UReal, [ URowVector ]); + add_nullary "e"; + add_unqualified ("eigenvalues_sym", ReturnType UVector, [ UMatrix ]); + add_unqualified ("eigenvectors_sym", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("qr_Q", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("qr_R", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("qr_thin_Q", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("qr_thin_R", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("elt_divide", ReturnType UVector, [ UVector; UVector ]); + add_unqualified ("elt_divide", ReturnType URowVector, [ URowVector; URowVector ]); + add_unqualified ("elt_divide", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("elt_divide", ReturnType UVector, [ UVector; UReal ]); + add_unqualified ("elt_divide", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("elt_divide", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("elt_divide", ReturnType UVector, [ UReal; UVector ]); + add_unqualified ("elt_divide", ReturnType URowVector, [ UReal; URowVector ]); + add_unqualified ("elt_divide", ReturnType UMatrix, [ UReal; UMatrix ]); + add_unqualified ("elt_multiply", ReturnType UVector, [ UVector; UVector ]); + add_unqualified ("elt_multiply", ReturnType URowVector, [ URowVector; URowVector ]); + add_unqualified ("elt_multiply", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("falling_factorial", ReturnType UReal, [ UReal; UInt ]); + add_unqualified ("falling_factorial", ReturnType UInt, [ UInt; UInt ]); + add_binary "fdim"; + add_ternary "fma"; + add_binary "fmax"; + add_binary "fmin"; + add_binary "fmod"; + add_binary "gamma_p"; + add_binary "gamma_q"; add_unqualified ( "gaussian_dlm_obs_log" , ReturnType UReal - , [UMatrix; UMatrix; UMatrix; UMatrix; UMatrix; UVector; UMatrix] ) ; + , [ UMatrix; UMatrix; UMatrix; UMatrix; UMatrix; UVector; UMatrix ] ); add_unqualified ( "gaussian_dlm_obs_log" , ReturnType UReal - , [UMatrix; UMatrix; UMatrix; UVector; UMatrix; UVector; UMatrix] ) ; + , [ UMatrix; UMatrix; UMatrix; UVector; UMatrix; UVector; UMatrix ] ); add_unqualified ( "gaussian_dlm_obs_lpdf" , ReturnType UReal - , [UMatrix; UMatrix; UMatrix; UMatrix; UMatrix; UVector; UMatrix] ) ; + , [ UMatrix; UMatrix; UMatrix; UMatrix; UMatrix; UVector; UMatrix ] ); add_unqualified ( "gaussian_dlm_obs_lpdf" , ReturnType UReal - , [UMatrix; UMatrix; UMatrix; UVector; UMatrix; UVector; UMatrix] ) ; + , [ UMatrix; UMatrix; UMatrix; UVector; UMatrix; UVector; UMatrix ] ); add_unqualified - ("gp_dot_prod_cov", ReturnType UMatrix, [bare_array_type (UReal, 1); UReal]) ; + ("gp_dot_prod_cov", ReturnType UMatrix, [ bare_array_type (UReal, 1); UReal ]); add_unqualified ( "gp_dot_prod_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal ] ); add_unqualified ( "gp_dot_prod_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal ] ); add_unqualified - ( "gp_dot_prod_cov" - , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal] ) ; + ("gp_dot_prod_cov", ReturnType UMatrix, [ bare_array_type (UVector, 1); UReal ]); add_unqualified ( "gp_dot_prod_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal] ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal ] ); add_unqualified - ( "gp_exp_quad_cov" - , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal] ) ; + ("gp_exp_quad_cov", ReturnType UMatrix, [ bare_array_type (UReal, 1); UReal; UReal ]); add_unqualified ( "gp_exp_quad_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified - ( "gp_exp_quad_cov" - , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal] ) ; + ("gp_exp_quad_cov", ReturnType UMatrix, [ bare_array_type (UVector, 1); UReal; UReal ]); add_unqualified ( "gp_exp_quad_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal] - ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "gp_exp_quad_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1)] ) ; + , [ bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1) ] ); add_unqualified ( "gp_exp_quad_cov" , ReturnType UMatrix , [ bare_array_type (UVector, 1) ; bare_array_type (UVector, 1) ; UReal - ; bare_array_type (UReal, 1) ] ) ; + ; bare_array_type (UReal, 1) + ] ); add_unqualified - ( "gp_matern32_cov" - , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal] ) ; + ("gp_matern32_cov", ReturnType UMatrix, [ bare_array_type (UReal, 1); UReal; UReal ]); add_unqualified ( "gp_matern32_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified - ( "gp_matern32_cov" - , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal] ) ; + ("gp_matern32_cov", ReturnType UMatrix, [ bare_array_type (UVector, 1); UReal; UReal ]); add_unqualified ( "gp_matern32_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal] - ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "gp_matern32_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1)] ) ; + , [ bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1) ] ); add_unqualified ( "gp_matern32_cov" , ReturnType UMatrix , [ bare_array_type (UVector, 1) ; bare_array_type (UVector, 1) ; UReal - ; bare_array_type (UReal, 1) ] ) ; + ; bare_array_type (UReal, 1) + ] ); add_unqualified - ( "gp_matern52_cov" - , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal] ) ; + ("gp_matern52_cov", ReturnType UMatrix, [ bare_array_type (UReal, 1); UReal; UReal ]); add_unqualified ( "gp_matern52_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified - ( "gp_matern52_cov" - , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal] ) ; + ("gp_matern52_cov", ReturnType UMatrix, [ bare_array_type (UVector, 1); UReal; UReal ]); add_unqualified ( "gp_matern52_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal] - ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "gp_matern52_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1)] ) ; + , [ bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1) ] ); add_unqualified ( "gp_matern52_cov" , ReturnType UMatrix , [ bare_array_type (UVector, 1) ; bare_array_type (UVector, 1) ; UReal - ; bare_array_type (UReal, 1) ] ) ; + ; bare_array_type (UReal, 1) + ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal] ) ; + , [ bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal] - ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1)] ) ; + , [ bare_array_type (UVector, 1); UReal; bare_array_type (UReal, 1) ] ); add_unqualified ( "gp_exponential_cov" , ReturnType UMatrix , [ bare_array_type (UVector, 1) ; bare_array_type (UVector, 1) ; UReal - ; bare_array_type (UReal, 1) ] ) ; + ; bare_array_type (UReal, 1) + ] ); add_unqualified ( "gp_periodic_cov" , ReturnType UMatrix - , [bare_array_type (UReal, 1); UReal; UReal; UReal] ) ; + , [ bare_array_type (UReal, 1); UReal; UReal; UReal ] ); add_unqualified ( "gp_periodic_cov" , ReturnType UMatrix - , [ bare_array_type (UReal, 1) - ; bare_array_type (UReal, 1) - ; UReal; UReal; UReal ] ) ; + , [ bare_array_type (UReal, 1); bare_array_type (UReal, 1); UReal; UReal; UReal ] ); add_unqualified ( "gp_periodic_cov" , ReturnType UMatrix - , [bare_array_type (UVector, 1); UReal; UReal; UReal] ) ; + , [ bare_array_type (UVector, 1); UReal; UReal; UReal ] ); add_unqualified ( "gp_periodic_cov" , ReturnType UMatrix - , [ bare_array_type (UVector, 1) - ; bare_array_type (UVector, 1) - ; UReal; UReal; UReal ] ) ; + , [ bare_array_type (UVector, 1); bare_array_type (UVector, 1); UReal; UReal; UReal ] + ); (* ; add_nullary ("get_lp") *) - add_unqualified ("head", ReturnType URowVector, [URowVector; UInt]) ; - add_unqualified ("head", ReturnType UVector, [UVector; UInt]) ; + add_unqualified ("head", ReturnType URowVector, [ URowVector; UInt ]); + add_unqualified ("head", ReturnType UVector, [ UVector; UInt ]); for i = 0 to bare_types_size - 1 do add_unqualified ( "head" , ReturnType (bare_array_type (bare_types i, 1)) - , [bare_array_type (bare_types i, 1); UInt] ) ; + , [ bare_array_type (bare_types i, 1); UInt ] ); add_unqualified ( "head" , ReturnType (bare_array_type (bare_types i, 2)) - , [bare_array_type (bare_types i, 2); UInt] ) ; + , [ bare_array_type (bare_types i, 2); UInt ] ); add_unqualified ( "head" , ReturnType (bare_array_type (bare_types i, 3)) - , [bare_array_type (bare_types i, 3); UInt] ) - done ; - add_unqualified - ("hmm_marginal", ReturnType UReal, [UMatrix; UMatrix; UVector]) ; + , [ bare_array_type (bare_types i, 3); UInt ] ) + done; + add_unqualified ("hmm_marginal", ReturnType UReal, [ UMatrix; UMatrix; UVector ]); add_qualified ( "hmm_hidden_state_prob" , ReturnType UMatrix - , [(DataOnly, UMatrix); (DataOnly, UMatrix); (DataOnly, UVector)] ) ; + , [ DataOnly, UMatrix; DataOnly, UMatrix; DataOnly, UVector ] ); add_unqualified ( "hmm_latent_rng" , ReturnType (bare_array_type (UInt, 1)) - , [UMatrix; UMatrix; UVector] ) ; - add_unqualified - ("hypergeometric_log", ReturnType UReal, [UInt; UInt; UInt; UInt]) ; - add_unqualified - ("hypergeometric_lpmf", ReturnType UReal, [UInt; UInt; UInt; UInt]) ; - add_unqualified ("hypergeometric_rng", ReturnType UInt, [UInt; UInt; UInt]) ; - add_binary "hypot" ; - add_unqualified ("identity_matrix", ReturnType UMatrix, [UInt]) ; + , [ UMatrix; UMatrix; UVector ] ); + add_unqualified ("hypergeometric_log", ReturnType UReal, [ UInt; UInt; UInt; UInt ]); + add_unqualified ("hypergeometric_lpmf", ReturnType UReal, [ UInt; UInt; UInt; UInt ]); + add_unqualified ("hypergeometric_rng", ReturnType UInt, [ UInt; UInt; UInt ]); + add_binary "hypot"; + add_unqualified ("identity_matrix", ReturnType UMatrix, [ UInt ]); for j = 0 to 8 - 1 do add_unqualified ( "if_else" , ReturnType (bare_array_type (UReal, j)) - , [UInt; bare_array_type (UReal, j); bare_array_type (UReal, j)] ) ; + , [ UInt; bare_array_type (UReal, j); bare_array_type (UReal, j) ] ); add_unqualified ( "if_else" , ReturnType (bare_array_type (UInt, j)) - , [UInt; bare_array_type (UInt, j); bare_array_type (UInt, j)] ) ; + , [ UInt; bare_array_type (UInt, j); bare_array_type (UInt, j) ] ); add_unqualified ( "if_else" , ReturnType (bare_array_type (UVector, j)) - , [UInt; bare_array_type (UVector, j); bare_array_type (UVector, j)] ) ; + , [ UInt; bare_array_type (UVector, j); bare_array_type (UVector, j) ] ); add_unqualified ( "if_else" , ReturnType (bare_array_type (URowVector, j)) - , [UInt; bare_array_type (URowVector, j); bare_array_type (URowVector, j)] - ) ; + , [ UInt; bare_array_type (URowVector, j); bare_array_type (URowVector, j) ] ); add_unqualified ( "if_else" , ReturnType (bare_array_type (UMatrix, j)) - , [UInt; bare_array_type (UMatrix, j); bare_array_type (UMatrix, j)] ) - done ; - add_unqualified ("inc_beta", ReturnType UReal, [UReal; UReal; UReal]) ; - add_unqualified ("int_step", ReturnType UInt, [UReal]) ; - add_unqualified ("int_step", ReturnType UInt, [UInt]) ; + , [ UInt; bare_array_type (UMatrix, j); bare_array_type (UMatrix, j) ] ) + done; + add_unqualified ("inc_beta", ReturnType UReal, [ UReal; UReal; UReal ]); + add_unqualified ("int_step", ReturnType UInt, [ UReal ]); + add_unqualified ("int_step", ReturnType UInt, [ UInt ]); add_qualified ( "integrate_1d" , ReturnType UReal , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal); (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UReal ) ) - ; (AutoDiffable, UReal); (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "integrate_1d" , ReturnType UReal , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal); (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UReal ) ) - ; (AutoDiffable, UReal); (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) ] - ) ; + ; AutoDiffable, UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ] ); add_qualified ( "integrate_ode" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "integrate_ode_adams" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "integrate_ode_adams" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) - ; (DataOnly, UReal); (DataOnly, UReal) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ; DataOnly, UReal + ; DataOnly, UReal + ] ); add_qualified ( "integrate_ode_bdf" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "integrate_ode_bdf" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) - ; (DataOnly, UReal); (DataOnly, UReal) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ; DataOnly, UReal + ; DataOnly, UReal + ] ); add_qualified ( "integrate_ode_rk45" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] ) ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] ); add_qualified ( "integrate_ode_rk45" , ReturnType (UArray (UArray UReal)) , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType (UArray UReal) ) ) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UReal) - ; (AutoDiffable, UArray UReal) - ; (AutoDiffable, UArray UReal) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt); (DataOnly, UReal) - ; (DataOnly, UReal); (DataOnly, UReal) ] ) ; - add_unqualified - ("inv_wishart_log", ReturnType UReal, [UMatrix; UReal; UMatrix]) ; - add_unqualified - ("inv_wishart_lpdf", ReturnType UReal, [UMatrix; UReal; UMatrix]) ; - add_unqualified ("inv_wishart_rng", ReturnType UMatrix, [UReal; UMatrix]) ; - add_unqualified ("inverse", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("inverse_spd", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("is_inf", ReturnType UInt, [UReal]) ; - add_unqualified ("is_nan", ReturnType UInt, [UReal]) ; - add_binary "lbeta" ; - add_binary "lchoose" ; - add_unqualified - ("linspaced_array", ReturnType (UArray UReal), [UInt; UReal; UReal]) ; - add_unqualified - ("linspaced_row_vector", ReturnType URowVector, [UInt; UReal; UReal]) ; - add_unqualified ("linspaced_vector", ReturnType UVector, [UInt; UReal; UReal]) ; - add_unqualified ("lkj_corr_cholesky_log", ReturnType UReal, [UMatrix; UReal]) ; - add_unqualified ("lkj_corr_cholesky_lpdf", ReturnType UReal, [UMatrix; UReal]) ; - add_unqualified ("lkj_corr_cholesky_rng", ReturnType UMatrix, [UInt; UReal]) ; - add_unqualified ("lkj_corr_log", ReturnType UReal, [UMatrix; UReal]) ; - add_unqualified ("lkj_corr_lpdf", ReturnType UReal, [UMatrix; UReal]) ; - add_unqualified ("lkj_corr_rng", ReturnType UMatrix, [UInt; UReal]) ; - add_unqualified - ("lkj_cov_log", ReturnType UReal, [UMatrix; UVector; UVector; UReal]) ; - add_unqualified ("lmgamma", ReturnType UReal, [UInt; UReal]) ; - add_binary "lmultiply" ; - add_nullary "log10" ; - add_nullary "log2" ; - add_unqualified ("log_determinant", ReturnType UReal, [UMatrix]) ; - add_binary "log_diff_exp" ; - add_binary "log_falling_factorial" ; - add_ternary "log_mix" ; + ; AutoDiffable, UArray UReal + ; AutoDiffable, UReal + ; AutoDiffable, UArray UReal + ; AutoDiffable, UArray UReal + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ; DataOnly, UReal + ; DataOnly, UReal + ; DataOnly, UReal + ] ); + add_unqualified ("inv_wishart_log", ReturnType UReal, [ UMatrix; UReal; UMatrix ]); + add_unqualified ("inv_wishart_lpdf", ReturnType UReal, [ UMatrix; UReal; UMatrix ]); + add_unqualified ("inv_wishart_rng", ReturnType UMatrix, [ UReal; UMatrix ]); + add_unqualified ("inverse", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("inverse_spd", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("is_inf", ReturnType UInt, [ UReal ]); + add_unqualified ("is_nan", ReturnType UInt, [ UReal ]); + add_binary "lbeta"; + add_binary "lchoose"; + add_unqualified ("linspaced_array", ReturnType (UArray UReal), [ UInt; UReal; UReal ]); + add_unqualified ("linspaced_row_vector", ReturnType URowVector, [ UInt; UReal; UReal ]); + add_unqualified ("linspaced_vector", ReturnType UVector, [ UInt; UReal; UReal ]); + add_unqualified ("lkj_corr_cholesky_log", ReturnType UReal, [ UMatrix; UReal ]); + add_unqualified ("lkj_corr_cholesky_lpdf", ReturnType UReal, [ UMatrix; UReal ]); + add_unqualified ("lkj_corr_cholesky_rng", ReturnType UMatrix, [ UInt; UReal ]); + add_unqualified ("lkj_corr_log", ReturnType UReal, [ UMatrix; UReal ]); + add_unqualified ("lkj_corr_lpdf", ReturnType UReal, [ UMatrix; UReal ]); + add_unqualified ("lkj_corr_rng", ReturnType UMatrix, [ UInt; UReal ]); + add_unqualified ("lkj_cov_log", ReturnType UReal, [ UMatrix; UVector; UVector; UReal ]); + add_unqualified ("lmgamma", ReturnType UReal, [ UInt; UReal ]); + add_binary "lmultiply"; + add_nullary "log10"; + add_nullary "log2"; + add_unqualified ("log_determinant", ReturnType UReal, [ UMatrix ]); + add_binary "log_diff_exp"; + add_binary "log_falling_factorial"; + add_ternary "log_mix"; for i = 1 to vector_types_size - 1 do for j = 1 to vector_types_size - 1 do - add_unqualified - ("log_mix", ReturnType UReal, [vector_types i; vector_types j]) - done ; + add_unqualified ("log_mix", ReturnType UReal, [ vector_types i; vector_types j ]) + done; add_unqualified - ( "log_mix" - , ReturnType UReal - , [vector_types i; bare_array_type (UVector, 1)] ) ; + ("log_mix", ReturnType UReal, [ vector_types i; bare_array_type (UVector, 1) ]); add_unqualified - ( "log_mix" - , ReturnType UReal - , [vector_types i; bare_array_type (URowVector, 1)] ) - done ; - add_binary "log_rising_factorial" ; - add_unqualified ("log_softmax", ReturnType UVector, [UVector]) ; - add_unqualified - ("log_sum_exp", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("log_sum_exp", ReturnType UReal, [UVector]) ; - add_unqualified ("log_sum_exp", ReturnType UReal, [URowVector]) ; - add_unqualified ("log_sum_exp", ReturnType UReal, [UMatrix]) ; - add_binary "log_sum_exp" ; + ("log_mix", ReturnType UReal, [ vector_types i; bare_array_type (URowVector, 1) ]) + done; + add_binary "log_rising_factorial"; + add_unqualified ("log_softmax", ReturnType UVector, [ UVector ]); + add_unqualified ("log_sum_exp", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("log_sum_exp", ReturnType UReal, [ UVector ]); + add_unqualified ("log_sum_exp", ReturnType UReal, [ URowVector ]); + add_unqualified ("log_sum_exp", ReturnType UReal, [ UMatrix ]); + add_binary "log_sum_exp"; for i = 0 to primitive_types_size - 1 do - add_unqualified ("logical_negation", ReturnType UInt, [primitive_types i]) ; + add_unqualified ("logical_negation", ReturnType UInt, [ primitive_types i ]); for j = 0 to primitive_types_size - 1 do add_unqualified - ("logical_or", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_or", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_and", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_and", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_eq", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_eq", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_neq", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_neq", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_lt", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_lt", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_lte", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_lte", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_gt", ReturnType UInt, [primitive_types i; primitive_types j]) ; + ("logical_gt", ReturnType UInt, [ primitive_types i; primitive_types j ]); add_unqualified - ("logical_gte", ReturnType UInt, [primitive_types i; primitive_types j]) + ("logical_gte", ReturnType UInt, [ primitive_types i; primitive_types j ]) done - done ; - add_nullary "machine_precision" ; + done; + add_nullary "machine_precision"; add_qualified ( "map_rect" , ReturnType UVector , [ ( AutoDiffable , UFun - ( [ (AutoDiffable, UVector); (AutoDiffable, UVector) - ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ] + ( [ AutoDiffable, UVector + ; AutoDiffable, UVector + ; DataOnly, UArray UReal + ; DataOnly, UArray UInt + ] , ReturnType UVector ) ) - ; (AutoDiffable, UVector) - ; (AutoDiffable, UArray UVector) - ; (DataOnly, UArray (UArray UReal)) - ; (DataOnly, UArray (UArray UInt)) ] ) ; - add_unqualified ("matrix_exp", ReturnType UMatrix, [UMatrix]) ; - add_unqualified - ("matrix_exp_multiply", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("max", ReturnType UInt, [bare_array_type (UInt, 1)]) ; - add_unqualified ("max", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("max", ReturnType UReal, [UVector]) ; - add_unqualified ("max", ReturnType UReal, [URowVector]) ; - add_unqualified ("max", ReturnType UReal, [UMatrix]) ; - add_unqualified ("max", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("mdivide_left", ReturnType UVector, [UMatrix; UVector]) ; - add_unqualified ("mdivide_left", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("mdivide_left_spd", ReturnType UVector, [UMatrix; UVector]) ; - add_unqualified ("mdivide_left_spd", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified - ("mdivide_left_tri_low", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified - ("mdivide_left_tri_low", ReturnType UVector, [UMatrix; UVector]) ; - add_unqualified - ("mdivide_right", ReturnType URowVector, [URowVector; UMatrix]) ; - add_unqualified ("mdivide_right_spd", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified - ("mdivide_right_spd", ReturnType URowVector, [URowVector; UMatrix]) ; - add_unqualified ("mdivide_right", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified - ("mdivide_right_tri_low", ReturnType URowVector, [URowVector; UMatrix]) ; - add_unqualified - ("mdivide_right_tri_low", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("mean", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("mean", ReturnType UReal, [UVector]) ; - add_unqualified ("mean", ReturnType UReal, [URowVector]) ; - add_unqualified ("mean", ReturnType UReal, [UMatrix]) ; - add_unqualified ("min", ReturnType UInt, [bare_array_type (UInt, 1)]) ; - add_unqualified ("min", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("min", ReturnType UReal, [UVector]) ; - add_unqualified ("min", ReturnType UReal, [URowVector]) ; - add_unqualified ("min", ReturnType UReal, [UMatrix]) ; - add_unqualified ("min", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("minus", ReturnType UInt, [UInt]) ; - add_unqualified ("minus", ReturnType UReal, [UReal]) ; - add_unqualified ("minus", ReturnType UVector, [UVector]) ; - add_unqualified ("minus", ReturnType URowVector, [URowVector]) ; - add_unqualified ("minus", ReturnType UMatrix, [UMatrix]) ; - add_unqualified - ("modified_bessel_first_kind", ReturnType UReal, [UInt; UReal]) ; - add_unqualified - ("modified_bessel_second_kind", ReturnType UReal, [UInt; UReal]) ; - add_unqualified ("modulus", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("multi_normal_rng", ReturnType UVector, [UVector; UMatrix]) ; + ; AutoDiffable, UVector + ; AutoDiffable, UArray UVector + ; DataOnly, UArray (UArray UReal) + ; DataOnly, UArray (UArray UInt) + ] ); + add_unqualified ("matrix_exp", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("matrix_exp_multiply", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("max", ReturnType UInt, [ bare_array_type (UInt, 1) ]); + add_unqualified ("max", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("max", ReturnType UReal, [ UVector ]); + add_unqualified ("max", ReturnType UReal, [ URowVector ]); + add_unqualified ("max", ReturnType UReal, [ UMatrix ]); + add_unqualified ("max", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("mdivide_left", ReturnType UVector, [ UMatrix; UVector ]); + add_unqualified ("mdivide_left", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mdivide_left_spd", ReturnType UVector, [ UMatrix; UVector ]); + add_unqualified ("mdivide_left_spd", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mdivide_left_tri_low", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mdivide_left_tri_low", ReturnType UVector, [ UMatrix; UVector ]); + add_unqualified ("mdivide_right", ReturnType URowVector, [ URowVector; UMatrix ]); + add_unqualified ("mdivide_right_spd", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mdivide_right_spd", ReturnType URowVector, [ URowVector; UMatrix ]); + add_unqualified ("mdivide_right", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mdivide_right_tri_low", ReturnType URowVector, [ URowVector; UMatrix ]); + add_unqualified ("mdivide_right_tri_low", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("mean", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("mean", ReturnType UReal, [ UVector ]); + add_unqualified ("mean", ReturnType UReal, [ URowVector ]); + add_unqualified ("mean", ReturnType UReal, [ UMatrix ]); + add_unqualified ("min", ReturnType UInt, [ bare_array_type (UInt, 1) ]); + add_unqualified ("min", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("min", ReturnType UReal, [ UVector ]); + add_unqualified ("min", ReturnType UReal, [ URowVector ]); + add_unqualified ("min", ReturnType UReal, [ UMatrix ]); + add_unqualified ("min", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("minus", ReturnType UInt, [ UInt ]); + add_unqualified ("minus", ReturnType UReal, [ UReal ]); + add_unqualified ("minus", ReturnType UVector, [ UVector ]); + add_unqualified ("minus", ReturnType URowVector, [ URowVector ]); + add_unqualified ("minus", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("modified_bessel_first_kind", ReturnType UReal, [ UInt; UReal ]); + add_unqualified ("modified_bessel_second_kind", ReturnType UReal, [ UInt; UReal ]); + add_unqualified ("modulus", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("multi_normal_rng", ReturnType UVector, [ UVector; UMatrix ]); add_unqualified ( "multi_normal_rng" , ReturnType (bare_array_type (UVector, 1)) - , [bare_array_type (UVector, 1); UMatrix] ) ; - add_unqualified - ("multi_normal_rng", ReturnType UVector, [URowVector; UMatrix]) ; + , [ bare_array_type (UVector, 1); UMatrix ] ); + add_unqualified ("multi_normal_rng", ReturnType UVector, [ URowVector; UMatrix ]); add_unqualified ( "multi_normal_rng" , ReturnType (bare_array_type (UVector, 1)) - , [bare_array_type (URowVector, 1); UMatrix] ) ; - add_unqualified - ("multi_normal_cholesky_rng", ReturnType UVector, [UVector; UMatrix]) ; + , [ bare_array_type (URowVector, 1); UMatrix ] ); + add_unqualified ("multi_normal_cholesky_rng", ReturnType UVector, [ UVector; UMatrix ]); add_unqualified ( "multi_normal_cholesky_rng" , ReturnType (bare_array_type (UVector, 1)) - , [bare_array_type (UVector, 1); UMatrix] ) ; + , [ bare_array_type (UVector, 1); UMatrix ] ); add_unqualified - ("multi_normal_cholesky_rng", ReturnType UVector, [URowVector; UMatrix]) ; + ("multi_normal_cholesky_rng", ReturnType UVector, [ URowVector; UMatrix ]); add_unqualified ( "multi_normal_cholesky_rng" , ReturnType (bare_array_type (UVector, 1)) - , [bare_array_type (URowVector, 1); UMatrix] ) ; - add_unqualified - ("multi_student_t_rng", ReturnType UVector, [UReal; UVector; UMatrix]) ; + , [ bare_array_type (URowVector, 1); UMatrix ] ); + add_unqualified ("multi_student_t_rng", ReturnType UVector, [ UReal; UVector; UMatrix ]); add_unqualified ( "multi_student_t_rng" , ReturnType (bare_array_type (UVector, 1)) - , [UReal; bare_array_type (UVector, 1); UMatrix] ) ; + , [ UReal; bare_array_type (UVector, 1); UMatrix ] ); add_unqualified - ("multi_student_t_rng", ReturnType UVector, [UReal; URowVector; UMatrix]) ; + ("multi_student_t_rng", ReturnType UVector, [ UReal; URowVector; UMatrix ]); add_unqualified ( "multi_student_t_rng" , ReturnType (bare_array_type (UVector, 1)) - , [UReal; bare_array_type (URowVector, 1); UMatrix] ) ; - add_unqualified - ("multinomial_log", ReturnType UReal, [bare_array_type (UInt, 1); UVector]) ; - add_unqualified - ("multinomial_lpmf", ReturnType UReal, [bare_array_type (UInt, 1); UVector]) ; - add_unqualified - ("multinomial_rng", ReturnType (bare_array_type (UInt, 1)), [UVector; UInt]) ; - add_unqualified ("multiply", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("multiply", ReturnType UReal, [UReal; UReal]) ; - add_unqualified ("multiply", ReturnType UVector, [UVector; UReal]) ; - add_unqualified ("multiply", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("multiply", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("multiply", ReturnType UReal, [URowVector; UVector]) ; - add_unqualified ("multiply", ReturnType UMatrix, [UVector; URowVector]) ; - add_unqualified ("multiply", ReturnType UVector, [UMatrix; UVector]) ; - add_unqualified ("multiply", ReturnType URowVector, [URowVector; UMatrix]) ; - add_unqualified ("multiply", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("multiply", ReturnType UVector, [UReal; UVector]) ; - add_unqualified ("multiply", ReturnType URowVector, [UReal; URowVector]) ; - add_unqualified ("multiply", ReturnType UMatrix, [UReal; UMatrix]) ; - add_binary "multiply_log" ; - add_unqualified - ("multiply_lower_tri_self_transpose", ReturnType UMatrix, [UMatrix]) ; + , [ UReal; bare_array_type (URowVector, 1); UMatrix ] ); + add_unqualified + ("multinomial_log", ReturnType UReal, [ bare_array_type (UInt, 1); UVector ]); + add_unqualified + ("multinomial_lpmf", ReturnType UReal, [ bare_array_type (UInt, 1); UVector ]); + add_unqualified + ("multinomial_rng", ReturnType (bare_array_type (UInt, 1)), [ UVector; UInt ]); + add_unqualified ("multiply", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("multiply", ReturnType UReal, [ UReal; UReal ]); + add_unqualified ("multiply", ReturnType UVector, [ UVector; UReal ]); + add_unqualified ("multiply", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("multiply", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("multiply", ReturnType UReal, [ URowVector; UVector ]); + add_unqualified ("multiply", ReturnType UMatrix, [ UVector; URowVector ]); + add_unqualified ("multiply", ReturnType UVector, [ UMatrix; UVector ]); + add_unqualified ("multiply", ReturnType URowVector, [ URowVector; UMatrix ]); + add_unqualified ("multiply", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("multiply", ReturnType UVector, [ UReal; UVector ]); + add_unqualified ("multiply", ReturnType URowVector, [ UReal; URowVector ]); + add_unqualified ("multiply", ReturnType UMatrix, [ UReal; UMatrix ]); + add_binary "multiply_log"; + add_unqualified ("multiply_lower_tri_self_transpose", ReturnType UMatrix, [ UMatrix ]); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UReal; UVector; UReal] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UReal; UVector; UReal ] ); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UVector; UVector; UReal] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UVector; UVector; UReal ] ); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [UInt; UMatrix; UReal; UVector; UReal] ) ; + , [ UInt; UMatrix; UReal; UVector; UReal ] ); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [UInt; UMatrix; UVector; UVector; UReal] ) ; + , [ UInt; UMatrix; UVector; UVector; UReal ] ); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UReal; UVector; UReal] ) ; + , [ bare_array_type (UInt, 1); URowVector; UReal; UVector; UReal ] ); add_unqualified ( "neg_binomial_2_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UVector; UVector; UReal] ) ; - add_nullary "negative_infinity" ; + , [ bare_array_type (UInt, 1); URowVector; UVector; UVector; UReal ] ); + add_nullary "negative_infinity"; add_unqualified - ( "normal_id_glm_lpdf" - , ReturnType UReal - , [UVector; UMatrix; UReal; UVector; UReal] ) ; + ("normal_id_glm_lpdf", ReturnType UReal, [ UVector; UMatrix; UReal; UVector; UReal ]); add_unqualified - ( "normal_id_glm_lpdf" - , ReturnType UReal - , [UVector; UMatrix; UVector; UVector; UReal] ) ; + ("normal_id_glm_lpdf", ReturnType UReal, [ UVector; UMatrix; UVector; UVector; UReal ]); add_unqualified - ( "normal_id_glm_lpdf" - , ReturnType UReal - , [UReal; UMatrix; UReal; UVector; UVector] ) ; + ("normal_id_glm_lpdf", ReturnType UReal, [ UReal; UMatrix; UReal; UVector; UVector ]); add_unqualified - ( "normal_id_glm_lpdf" - , ReturnType UReal - , [UReal; UMatrix; UVector; UVector; UVector] ) ; + ("normal_id_glm_lpdf", ReturnType UReal, [ UReal; UMatrix; UVector; UVector; UVector ]); add_unqualified ( "normal_id_glm_lpdf" , ReturnType UReal - , [UVector; URowVector; UReal; UVector; UVector] ) ; + , [ UVector; URowVector; UReal; UVector; UVector ] ); add_unqualified ( "normal_id_glm_lpdf" , ReturnType UReal - , [UVector; URowVector; UVector; UVector; UVector] ) ; - add_nullary "not_a_number" ; - add_unqualified ("num_elements", ReturnType UInt, [UMatrix]) ; - add_unqualified ("num_elements", ReturnType UInt, [UVector]) ; - add_unqualified ("num_elements", ReturnType UInt, [URowVector]) ; + , [ UVector; URowVector; UVector; UVector; UVector ] ); + add_nullary "not_a_number"; + add_unqualified ("num_elements", ReturnType UInt, [ UMatrix ]); + add_unqualified ("num_elements", ReturnType UInt, [ UVector ]); + add_unqualified ("num_elements", ReturnType UInt, [ URowVector ]); for i = 1 to 10 - 1 do - add_unqualified - ("num_elements", ReturnType UInt, [bare_array_type (UInt, i)]) ; - add_unqualified - ("num_elements", ReturnType UInt, [bare_array_type (UReal, i)]) ; - add_unqualified - ("num_elements", ReturnType UInt, [bare_array_type (UMatrix, i)]) ; - add_unqualified - ("num_elements", ReturnType UInt, [bare_array_type (URowVector, i)]) ; - add_unqualified - ("num_elements", ReturnType UInt, [bare_array_type (UVector, i)]) - done ; - add_unqualified ("one_hot_int_array", ReturnType (UArray UInt), [UInt; UInt]) ; - add_unqualified ("one_hot_array", ReturnType (UArray UReal), [UInt; UInt]) ; - add_unqualified ("one_hot_row_vector", ReturnType URowVector, [UInt; UInt]) ; - add_unqualified ("one_hot_vector", ReturnType UVector, [UInt; UInt]) ; - add_unqualified ("ones_int_array", ReturnType (UArray UInt), [UInt]) ; - add_unqualified ("ones_array", ReturnType (UArray UReal), [UInt]) ; - add_unqualified ("ones_row_vector", ReturnType URowVector, [UInt]) ; - add_unqualified ("ones_vector", ReturnType UVector, [UInt]) ; + add_unqualified ("num_elements", ReturnType UInt, [ bare_array_type (UInt, i) ]); + add_unqualified ("num_elements", ReturnType UInt, [ bare_array_type (UReal, i) ]); + add_unqualified ("num_elements", ReturnType UInt, [ bare_array_type (UMatrix, i) ]); + add_unqualified ("num_elements", ReturnType UInt, [ bare_array_type (URowVector, i) ]); + add_unqualified ("num_elements", ReturnType UInt, [ bare_array_type (UVector, i) ]) + done; + add_unqualified ("one_hot_int_array", ReturnType (UArray UInt), [ UInt; UInt ]); + add_unqualified ("one_hot_array", ReturnType (UArray UReal), [ UInt; UInt ]); + add_unqualified ("one_hot_row_vector", ReturnType URowVector, [ UInt; UInt ]); + add_unqualified ("one_hot_vector", ReturnType UVector, [ UInt; UInt ]); + add_unqualified ("ones_int_array", ReturnType (UArray UInt), [ UInt ]); + add_unqualified ("ones_array", ReturnType (UArray UReal), [ UInt ]); + add_unqualified ("ones_row_vector", ReturnType URowVector, [ UInt ]); + add_unqualified ("ones_vector", ReturnType UVector, [ UInt ]); add_unqualified ( "ordered_logistic_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UVector; UVector ] ); add_unqualified ( "ordered_logistic_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); URowVector; UVector; UVector ] ); add_unqualified - ( "ordered_logistic_glm_lpmf" - , ReturnType UReal - , [UInt; UMatrix; UVector; UVector] ) ; - add_unqualified - ( "ordered_logistic_glm_lpmf" - , ReturnType UReal - , [UInt; URowVector; UVector; UVector] ) ; + ("ordered_logistic_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UVector; UVector ]); add_unqualified - ("ordered_logistic_log", ReturnType UReal, [UInt; UReal; UVector]) ; + ("ordered_logistic_glm_lpmf", ReturnType UReal, [ UInt; URowVector; UVector; UVector ]); + add_unqualified ("ordered_logistic_log", ReturnType UReal, [ UInt; UReal; UVector ]); add_unqualified ( "ordered_logistic_log" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UVector; UVector ] ); add_unqualified ( "ordered_logistic_log" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1)] ) ; - add_unqualified - ("ordered_logistic_lpmf", ReturnType UReal, [UInt; UReal; UVector]) ; + , [ bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1) ] ); + add_unqualified ("ordered_logistic_lpmf", ReturnType UReal, [ UInt; UReal; UVector ]); add_unqualified ( "ordered_logistic_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UVector; UVector ] ); add_unqualified ( "ordered_logistic_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1)] ) ; - add_unqualified ("ordered_logistic_rng", ReturnType UInt, [UReal; UVector]) ; - add_unqualified - ("ordered_probit_log", ReturnType UReal, [UInt; UReal; UVector]) ; + , [ bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1) ] ); + add_unqualified ("ordered_logistic_rng", ReturnType UInt, [ UReal; UVector ]); + add_unqualified ("ordered_probit_log", ReturnType UReal, [ UInt; UReal; UVector ]); add_unqualified ( "ordered_probit_log" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UVector; UVector ] ); add_unqualified ( "ordered_probit_log" , ReturnType UReal - , [bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1)] ) ; - add_unqualified - ("ordered_probit_lpmf", ReturnType UReal, [UInt; UReal; UVector]) ; + , [ bare_array_type (UInt, 1); UVector; bare_array_type (UVector, 1) ] ); + add_unqualified ("ordered_probit_lpmf", ReturnType UReal, [ UInt; UReal; UVector ]); add_unqualified ( "ordered_probit_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UReal; UVector] ) ; + , [ bare_array_type (UInt, 1); UReal; UVector ] ); add_unqualified ( "ordered_probit_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UReal; bare_array_type (UVector, 1)] ) ; - add_unqualified ("ordered_probit_rng", ReturnType UInt, [UReal; UVector]) ; - add_binary "owens_t" ; - add_nullary "pi" ; - add_unqualified ("plus", ReturnType UInt, [UInt]) ; - add_unqualified ("plus", ReturnType UReal, [UReal]) ; - add_unqualified ("plus", ReturnType UVector, [UVector]) ; - add_unqualified ("plus", ReturnType URowVector, [URowVector]) ; - add_unqualified ("plus", ReturnType UMatrix, [UMatrix]) ; + , [ bare_array_type (UInt, 1); UReal; bare_array_type (UVector, 1) ] ); + add_unqualified ("ordered_probit_rng", ReturnType UInt, [ UReal; UVector ]); + add_binary "owens_t"; + add_nullary "pi"; + add_unqualified ("plus", ReturnType UInt, [ UInt ]); + add_unqualified ("plus", ReturnType UReal, [ UReal ]); + add_unqualified ("plus", ReturnType UVector, [ UVector ]); + add_unqualified ("plus", ReturnType URowVector, [ URowVector ]); + add_unqualified ("plus", ReturnType UMatrix, [ UMatrix ]); add_unqualified ( "poisson_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UReal; UVector] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UReal; UVector ] ); add_unqualified ( "poisson_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); UMatrix; UVector; UVector] ) ; + , [ bare_array_type (UInt, 1); UMatrix; UVector; UVector ] ); add_unqualified - ("poisson_log_glm_lpmf", ReturnType UReal, [UInt; UMatrix; UReal; UVector]) ; + ("poisson_log_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UReal; UVector ]); add_unqualified - ( "poisson_log_glm_lpmf" - , ReturnType UReal - , [UInt; UMatrix; UVector; UVector] ) ; + ("poisson_log_glm_lpmf", ReturnType UReal, [ UInt; UMatrix; UVector; UVector ]); add_unqualified ( "poisson_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UReal; UVector] ) ; + , [ bare_array_type (UInt, 1); URowVector; UReal; UVector ] ); add_unqualified ( "poisson_log_glm_lpmf" , ReturnType UReal - , [bare_array_type (UInt, 1); URowVector; UVector; UVector] ) ; - add_nullary "positive_infinity" ; - add_binary "pow" ; - add_unqualified ("prod", ReturnType UInt, [bare_array_type (UInt, 1)]) ; - add_unqualified ("prod", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("prod", ReturnType UReal, [UVector]) ; - add_unqualified ("prod", ReturnType UReal, [URowVector]) ; - add_unqualified ("prod", ReturnType UReal, [UMatrix]) ; - add_unqualified ("quad_form", ReturnType UReal, [UMatrix; UVector]) ; - add_unqualified ("quad_form", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("quad_form_sym", ReturnType UReal, [UMatrix; UVector]) ; - add_unqualified ("quad_form_sym", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("quad_form_diag", ReturnType UMatrix, [UMatrix; UVector]) ; - add_unqualified ("quad_form_diag", ReturnType UMatrix, [UMatrix; URowVector]) ; - add_unqualified ("rank", ReturnType UInt, [bare_array_type (UInt, 1); UInt]) ; - add_unqualified ("rank", ReturnType UInt, [bare_array_type (UReal, 1); UInt]) ; - add_unqualified ("rank", ReturnType UInt, [UVector; UInt]) ; - add_unqualified ("rank", ReturnType UInt, [URowVector; UInt]) ; - add_unqualified ("append_row", ReturnType UMatrix, [UMatrix; UMatrix]) ; - add_unqualified ("append_row", ReturnType UMatrix, [URowVector; UMatrix]) ; - add_unqualified ("append_row", ReturnType UMatrix, [UMatrix; URowVector]) ; - add_unqualified ("append_row", ReturnType UMatrix, [URowVector; URowVector]) ; - add_unqualified ("append_row", ReturnType UVector, [UVector; UVector]) ; - add_unqualified ("append_row", ReturnType UVector, [UReal; UVector]) ; - add_unqualified ("append_row", ReturnType UVector, [UVector; UReal]) ; + , [ bare_array_type (UInt, 1); URowVector; UVector; UVector ] ); + add_nullary "positive_infinity"; + add_binary "pow"; + add_unqualified ("prod", ReturnType UInt, [ bare_array_type (UInt, 1) ]); + add_unqualified ("prod", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("prod", ReturnType UReal, [ UVector ]); + add_unqualified ("prod", ReturnType UReal, [ URowVector ]); + add_unqualified ("prod", ReturnType UReal, [ UMatrix ]); + add_unqualified ("quad_form", ReturnType UReal, [ UMatrix; UVector ]); + add_unqualified ("quad_form", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("quad_form_sym", ReturnType UReal, [ UMatrix; UVector ]); + add_unqualified ("quad_form_sym", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("quad_form_diag", ReturnType UMatrix, [ UMatrix; UVector ]); + add_unqualified ("quad_form_diag", ReturnType UMatrix, [ UMatrix; URowVector ]); + add_unqualified ("rank", ReturnType UInt, [ bare_array_type (UInt, 1); UInt ]); + add_unqualified ("rank", ReturnType UInt, [ bare_array_type (UReal, 1); UInt ]); + add_unqualified ("rank", ReturnType UInt, [ UVector; UInt ]); + add_unqualified ("rank", ReturnType UInt, [ URowVector; UInt ]); + add_unqualified ("append_row", ReturnType UMatrix, [ UMatrix; UMatrix ]); + add_unqualified ("append_row", ReturnType UMatrix, [ URowVector; UMatrix ]); + add_unqualified ("append_row", ReturnType UMatrix, [ UMatrix; URowVector ]); + add_unqualified ("append_row", ReturnType UMatrix, [ URowVector; URowVector ]); + add_unqualified ("append_row", ReturnType UVector, [ UVector; UVector ]); + add_unqualified ("append_row", ReturnType UVector, [ UReal; UVector ]); + add_unqualified ("append_row", ReturnType UVector, [ UVector; UReal ]); for i = 0 to bare_types_size - 1 do add_unqualified - ( "rep_array" - , ReturnType (bare_array_type (bare_types i, 1)) - , [bare_types i; UInt] ) ; + ("rep_array", ReturnType (bare_array_type (bare_types i, 1)), [ bare_types i; UInt ]); add_unqualified ( "rep_array" , ReturnType (bare_array_type (bare_types i, 2)) - , [bare_types i; UInt; UInt] ) ; + , [ bare_types i; UInt; UInt ] ); add_unqualified ( "rep_array" , ReturnType (bare_array_type (bare_types i, 3)) - , [bare_types i; UInt; UInt; UInt] ) ; + , [ bare_types i; UInt; UInt; UInt ] ); for j = 1 to 3 - 1 do add_unqualified ( "rep_array" , ReturnType (bare_array_type (bare_types i, j + 1)) - , [bare_array_type (bare_types i, j); UInt] ) ; + , [ bare_array_type (bare_types i, j); UInt ] ); add_unqualified ( "rep_array" , ReturnType (bare_array_type (bare_types i, j + 2)) - , [bare_array_type (bare_types i, j); UInt; UInt] ) ; + , [ bare_array_type (bare_types i, j); UInt; UInt ] ); add_unqualified ( "rep_array" , ReturnType (bare_array_type (bare_types i, j + 3)) - , [bare_array_type (bare_types i, j); UInt; UInt; UInt] ) + , [ bare_array_type (bare_types i, j); UInt; UInt; UInt ] ) done - done ; - add_unqualified ("rep_matrix", ReturnType UMatrix, [UReal; UInt; UInt]) ; - add_unqualified ("rep_matrix", ReturnType UMatrix, [UVector; UInt]) ; - add_unqualified ("rep_matrix", ReturnType UMatrix, [URowVector; UInt]) ; - add_unqualified ("rep_row_vector", ReturnType URowVector, [UReal; UInt]) ; - add_unqualified ("rep_vector", ReturnType UVector, [UReal; UInt]) ; + done; + add_unqualified ("rep_matrix", ReturnType UMatrix, [ UReal; UInt; UInt ]); + add_unqualified ("rep_matrix", ReturnType UMatrix, [ UVector; UInt ]); + add_unqualified ("rep_matrix", ReturnType UMatrix, [ URowVector; UInt ]); + add_unqualified ("rep_row_vector", ReturnType URowVector, [ UReal; UInt ]); + add_unqualified ("rep_vector", ReturnType UVector, [ UReal; UInt ]); for i = 0 to 7 do add_unqualified ( "reverse" , ReturnType (bare_array_type (UVector, i)) - , [bare_array_type (UVector, i)] ) ; + , [ bare_array_type (UVector, i) ] ); add_unqualified ( "reverse" , ReturnType (bare_array_type (URowVector, i)) - , [bare_array_type (URowVector, i)] ) - done ; + , [ bare_array_type (URowVector, i) ] ) + done; for i = 1 to 7 do add_unqualified - ( "reverse" - , ReturnType (bare_array_type (UInt, i)) - , [bare_array_type (UInt, i)] ) ; + ("reverse", ReturnType (bare_array_type (UInt, i)), [ bare_array_type (UInt, i) ]); add_unqualified - ( "reverse" - , ReturnType (bare_array_type (UReal, i)) - , [bare_array_type (UReal, i)] ) ; + ("reverse", ReturnType (bare_array_type (UReal, i)), [ bare_array_type (UReal, i) ]); add_unqualified ( "reverse" , ReturnType (bare_array_type (UMatrix, i)) - , [bare_array_type (UMatrix, i)] ) - done ; - add_unqualified ("rising_factorial", ReturnType UReal, [UReal; UInt]) ; - add_unqualified ("rising_factorial", ReturnType UInt, [UInt; UInt]) ; - add_unqualified ("row", ReturnType URowVector, [UMatrix; UInt]) ; - add_unqualified ("rows", ReturnType UInt, [UVector]) ; - add_unqualified ("rows", ReturnType UInt, [URowVector]) ; - add_unqualified ("rows", ReturnType UInt, [UMatrix]) ; - add_unqualified ("rows_dot_product", ReturnType UVector, [UVector; UVector]) ; - add_unqualified - ("rows_dot_product", ReturnType UVector, [URowVector; URowVector]) ; - add_unqualified ("rows_dot_product", ReturnType UVector, [UMatrix; UMatrix]) ; - add_unqualified ("rows_dot_self", ReturnType UVector, [UVector]) ; - add_unqualified ("rows_dot_self", ReturnType UVector, [URowVector]) ; - add_unqualified ("rows_dot_self", ReturnType UVector, [UMatrix]) ; - add_unqualified - ("scale_matrix_exp_multiply", ReturnType UMatrix, [UReal; UMatrix; UMatrix]) ; - add_unqualified ("sd", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("sd", ReturnType UReal, [UVector]) ; - add_unqualified ("sd", ReturnType UReal, [URowVector]) ; - add_unqualified ("sd", ReturnType UReal, [UMatrix]) ; - add_unqualified ("segment", ReturnType URowVector, [URowVector; UInt; UInt]) ; - add_unqualified ("segment", ReturnType UVector, [UVector; UInt; UInt]) ; + , [ bare_array_type (UMatrix, i) ] ) + done; + add_unqualified ("rising_factorial", ReturnType UReal, [ UReal; UInt ]); + add_unqualified ("rising_factorial", ReturnType UInt, [ UInt; UInt ]); + add_unqualified ("row", ReturnType URowVector, [ UMatrix; UInt ]); + add_unqualified ("rows", ReturnType UInt, [ UVector ]); + add_unqualified ("rows", ReturnType UInt, [ URowVector ]); + add_unqualified ("rows", ReturnType UInt, [ UMatrix ]); + add_unqualified ("rows_dot_product", ReturnType UVector, [ UVector; UVector ]); + add_unqualified ("rows_dot_product", ReturnType UVector, [ URowVector; URowVector ]); + add_unqualified ("rows_dot_product", ReturnType UVector, [ UMatrix; UMatrix ]); + add_unqualified ("rows_dot_self", ReturnType UVector, [ UVector ]); + add_unqualified ("rows_dot_self", ReturnType UVector, [ URowVector ]); + add_unqualified ("rows_dot_self", ReturnType UVector, [ UMatrix ]); + add_unqualified + ("scale_matrix_exp_multiply", ReturnType UMatrix, [ UReal; UMatrix; UMatrix ]); + add_unqualified ("sd", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("sd", ReturnType UReal, [ UVector ]); + add_unqualified ("sd", ReturnType UReal, [ URowVector ]); + add_unqualified ("sd", ReturnType UReal, [ UMatrix ]); + add_unqualified ("segment", ReturnType URowVector, [ URowVector; UInt; UInt ]); + add_unqualified ("segment", ReturnType UVector, [ UVector; UInt; UInt ]); for i = 0 to bare_types_size - 1 do add_unqualified ( "segment" , ReturnType (bare_array_type (bare_types i, 1)) - , [bare_array_type (bare_types i, 1); UInt; UInt] ) ; + , [ bare_array_type (bare_types i, 1); UInt; UInt ] ); add_unqualified ( "segment" , ReturnType (bare_array_type (bare_types i, 2)) - , [bare_array_type (bare_types i, 2); UInt; UInt] ) ; + , [ bare_array_type (bare_types i, 2); UInt; UInt ] ); add_unqualified ( "segment" , ReturnType (bare_array_type (bare_types i, 3)) - , [bare_array_type (bare_types i, 3); UInt; UInt] ) - done ; - add_unqualified ("singular_values", ReturnType UVector, [UMatrix]) ; + , [ bare_array_type (bare_types i, 3); UInt; UInt ] ) + done; + add_unqualified ("singular_values", ReturnType UVector, [ UMatrix ]); for i = 1 to 8 - 1 do - add_unqualified ("size", ReturnType UInt, [bare_array_type (UInt, i)]) ; - add_unqualified ("size", ReturnType UInt, [bare_array_type (UReal, i)]) ; - add_unqualified ("size", ReturnType UInt, [bare_array_type (UVector, i)]) ; - add_unqualified ("size", ReturnType UInt, [bare_array_type (URowVector, i)]) ; - add_unqualified ("size", ReturnType UInt, [bare_array_type (UMatrix, i)]) - done ; + add_unqualified ("size", ReturnType UInt, [ bare_array_type (UInt, i) ]); + add_unqualified ("size", ReturnType UInt, [ bare_array_type (UReal, i) ]); + add_unqualified ("size", ReturnType UInt, [ bare_array_type (UVector, i) ]); + add_unqualified ("size", ReturnType UInt, [ bare_array_type (URowVector, i) ]); + add_unqualified ("size", ReturnType UInt, [ bare_array_type (UMatrix, i) ]) + done; for i = 0 to bare_types_size - 1 do - add_unqualified ("size", ReturnType UInt, [bare_types i]) - done ; - add_unqualified ("softmax", ReturnType UVector, [UVector]) ; + add_unqualified ("size", ReturnType UInt, [ bare_types i ]) + done; + add_unqualified ("softmax", ReturnType UVector, [ UVector ]); add_unqualified - ( "sort_asc" - , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, 1)] ) ; + ("sort_asc", ReturnType (bare_array_type (UInt, 1)), [ bare_array_type (UInt, 1) ]); add_unqualified - ( "sort_asc" - , ReturnType (bare_array_type (UReal, 1)) - , [bare_array_type (UReal, 1)] ) ; - add_unqualified ("sort_asc", ReturnType UVector, [UVector]) ; - add_unqualified ("sort_asc", ReturnType URowVector, [URowVector]) ; + ("sort_asc", ReturnType (bare_array_type (UReal, 1)), [ bare_array_type (UReal, 1) ]); + add_unqualified ("sort_asc", ReturnType UVector, [ UVector ]); + add_unqualified ("sort_asc", ReturnType URowVector, [ URowVector ]); add_unqualified - ( "sort_desc" - , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, 1)] ) ; + ("sort_desc", ReturnType (bare_array_type (UInt, 1)), [ bare_array_type (UInt, 1) ]); add_unqualified - ( "sort_desc" - , ReturnType (bare_array_type (UReal, 1)) - , [bare_array_type (UReal, 1)] ) ; - add_unqualified ("sort_desc", ReturnType UVector, [UVector]) ; - add_unqualified ("sort_desc", ReturnType URowVector, [URowVector]) ; + ("sort_desc", ReturnType (bare_array_type (UReal, 1)), [ bare_array_type (UReal, 1) ]); + add_unqualified ("sort_desc", ReturnType UVector, [ UVector ]); + add_unqualified ("sort_desc", ReturnType URowVector, [ URowVector ]); add_unqualified ( "sort_indices_asc" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, 1)] ) ; + , [ bare_array_type (UInt, 1) ] ); add_unqualified ( "sort_indices_asc" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UReal, 1)] ) ; - add_unqualified - ("sort_indices_asc", ReturnType (bare_array_type (UInt, 1)), [UVector]) ; + , [ bare_array_type (UReal, 1) ] ); + add_unqualified ("sort_indices_asc", ReturnType (bare_array_type (UInt, 1)), [ UVector ]); add_unqualified - ("sort_indices_asc", ReturnType (bare_array_type (UInt, 1)), [URowVector]) ; + ("sort_indices_asc", ReturnType (bare_array_type (UInt, 1)), [ URowVector ]); add_unqualified ( "sort_indices_desc" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, 1)] ) ; + , [ bare_array_type (UInt, 1) ] ); add_unqualified ( "sort_indices_desc" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UReal, 1)] ) ; - add_unqualified - ("sort_indices_desc", ReturnType (bare_array_type (UInt, 1)), [UVector]) ; - add_unqualified - ("sort_indices_desc", ReturnType (bare_array_type (UInt, 1)), [URowVector]) ; - add_unqualified ("squared_distance", ReturnType UReal, [UReal; UReal]) ; - add_unqualified ("squared_distance", ReturnType UReal, [UVector; UVector]) ; - add_unqualified - ("squared_distance", ReturnType UReal, [URowVector; URowVector]) ; - add_unqualified ("squared_distance", ReturnType UReal, [UVector; URowVector]) ; - add_unqualified ("squared_distance", ReturnType UReal, [URowVector; UVector]) ; - add_nullary "sqrt2" ; - add_unqualified ("sub_col", ReturnType UVector, [UMatrix; UInt; UInt; UInt]) ; - add_unqualified - ("sub_row", ReturnType URowVector, [UMatrix; UInt; UInt; UInt]) ; + , [ bare_array_type (UReal, 1) ] ); + add_unqualified + ("sort_indices_desc", ReturnType (bare_array_type (UInt, 1)), [ UVector ]); + add_unqualified + ("sort_indices_desc", ReturnType (bare_array_type (UInt, 1)), [ URowVector ]); + add_unqualified ("squared_distance", ReturnType UReal, [ UReal; UReal ]); + add_unqualified ("squared_distance", ReturnType UReal, [ UVector; UVector ]); + add_unqualified ("squared_distance", ReturnType UReal, [ URowVector; URowVector ]); + add_unqualified ("squared_distance", ReturnType UReal, [ UVector; URowVector ]); + add_unqualified ("squared_distance", ReturnType UReal, [ URowVector; UVector ]); + add_nullary "sqrt2"; + add_unqualified ("sub_col", ReturnType UVector, [ UMatrix; UInt; UInt; UInt ]); + add_unqualified ("sub_row", ReturnType URowVector, [ UMatrix; UInt; UInt; UInt ]); for i = 0 to bare_types_size - 1 do - add_unqualified - ("subtract", ReturnType (bare_types i), [bare_types i; bare_types i]) - done ; - add_unqualified ("subtract", ReturnType UVector, [UVector; UReal]) ; - add_unqualified ("subtract", ReturnType URowVector, [URowVector; UReal]) ; - add_unqualified ("subtract", ReturnType UMatrix, [UMatrix; UReal]) ; - add_unqualified ("subtract", ReturnType UVector, [UReal; UVector]) ; - add_unqualified ("subtract", ReturnType URowVector, [UReal; URowVector]) ; - add_unqualified ("subtract", ReturnType UMatrix, [UReal; UMatrix]) ; - add_unqualified ("sum", ReturnType UInt, [bare_array_type (UInt, 1)]) ; - add_unqualified ("sum", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("sum", ReturnType UReal, [UVector]) ; - add_unqualified ("sum", ReturnType UReal, [URowVector]) ; - add_unqualified ("sum", ReturnType UReal, [UMatrix]) ; - add_unqualified ("tail", ReturnType URowVector, [URowVector; UInt]) ; - add_unqualified ("tail", ReturnType UVector, [UVector; UInt]) ; + add_unqualified ("subtract", ReturnType (bare_types i), [ bare_types i; bare_types i ]) + done; + add_unqualified ("subtract", ReturnType UVector, [ UVector; UReal ]); + add_unqualified ("subtract", ReturnType URowVector, [ URowVector; UReal ]); + add_unqualified ("subtract", ReturnType UMatrix, [ UMatrix; UReal ]); + add_unqualified ("subtract", ReturnType UVector, [ UReal; UVector ]); + add_unqualified ("subtract", ReturnType URowVector, [ UReal; URowVector ]); + add_unqualified ("subtract", ReturnType UMatrix, [ UReal; UMatrix ]); + add_unqualified ("sum", ReturnType UInt, [ bare_array_type (UInt, 1) ]); + add_unqualified ("sum", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("sum", ReturnType UReal, [ UVector ]); + add_unqualified ("sum", ReturnType UReal, [ URowVector ]); + add_unqualified ("sum", ReturnType UReal, [ UMatrix ]); + add_unqualified ("tail", ReturnType URowVector, [ URowVector; UInt ]); + add_unqualified ("tail", ReturnType UVector, [ UVector; UInt ]); for i = 0 to bare_types_size - 1 do add_unqualified ( "tail" , ReturnType (bare_array_type (bare_types i, 1)) - , [bare_array_type (bare_types i, 1); UInt] ) ; + , [ bare_array_type (bare_types i, 1); UInt ] ); add_unqualified ( "tail" , ReturnType (bare_array_type (bare_types i, 2)) - , [bare_array_type (bare_types i, 2); UInt] ) ; + , [ bare_array_type (bare_types i, 2); UInt ] ); add_unqualified ( "tail" , ReturnType (bare_array_type (bare_types i, 3)) - , [bare_array_type (bare_types i, 3); UInt] ) - done ; - add_unqualified ("tcrossprod", ReturnType UMatrix, [UMatrix]) ; - add_unqualified - ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [UMatrix]) ; - add_unqualified - ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [UVector]) ; - add_unqualified - ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [URowVector]) ; + , [ bare_array_type (bare_types i, 3); UInt ] ) + done; + add_unqualified ("tcrossprod", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [ UMatrix ]); + add_unqualified ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [ UVector ]); + add_unqualified ("to_array_1d", ReturnType (bare_array_type (UReal, 1)), [ URowVector ]); for i = 1 to 10 - 1 do add_unqualified ( "to_array_1d" , ReturnType (bare_array_type (UReal, 1)) - , [bare_array_type (UReal, i)] ) ; + , [ bare_array_type (UReal, i) ] ); add_unqualified ( "to_array_1d" , ReturnType (bare_array_type (UInt, 1)) - , [bare_array_type (UInt, i)] ) - done ; - add_unqualified - ("to_array_2d", ReturnType (bare_array_type (UReal, 2)), [UMatrix]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UMatrix; UInt; UInt]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UMatrix; UInt; UInt; UInt]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UVector]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UVector; UInt; UInt]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [UVector; UInt; UInt; UInt]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [URowVector]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [URowVector; UInt; UInt]) ; - add_unqualified - ("to_matrix", ReturnType UMatrix, [URowVector; UInt; UInt; UInt]) ; - add_unqualified - ("to_matrix", ReturnType UMatrix, [bare_array_type (UReal, 1); UInt; UInt]) ; - add_unqualified - ( "to_matrix" - , ReturnType UMatrix - , [bare_array_type (UReal, 1); UInt; UInt; UInt] ) ; - add_unqualified - ("to_matrix", ReturnType UMatrix, [bare_array_type (UInt, 1); UInt; UInt]) ; - add_unqualified - ( "to_matrix" - , ReturnType UMatrix - , [bare_array_type (UInt, 1); UInt; UInt; UInt] ) ; - add_unqualified - ("to_matrix", ReturnType UMatrix, [bare_array_type (UReal, 2)]) ; - add_unqualified ("to_matrix", ReturnType UMatrix, [bare_array_type (UInt, 2)]) ; - add_unqualified ("to_row_vector", ReturnType URowVector, [UMatrix]) ; - add_unqualified ("to_row_vector", ReturnType URowVector, [UVector]) ; - add_unqualified ("to_row_vector", ReturnType URowVector, [URowVector]) ; - add_unqualified - ("to_row_vector", ReturnType URowVector, [bare_array_type (UReal, 1)]) ; - add_unqualified - ("to_row_vector", ReturnType URowVector, [bare_array_type (UInt, 1)]) ; - add_unqualified ("to_vector", ReturnType UVector, [UMatrix]) ; - add_unqualified ("to_vector", ReturnType UVector, [UVector]) ; - add_unqualified ("to_vector", ReturnType UVector, [URowVector]) ; - add_unqualified - ("to_vector", ReturnType UVector, [bare_array_type (UReal, 1)]) ; - add_unqualified ("to_vector", ReturnType UVector, [bare_array_type (UInt, 1)]) ; - add_unqualified ("trace", ReturnType UReal, [UMatrix]) ; - add_unqualified - ("trace_gen_quad_form", ReturnType UReal, [UMatrix; UMatrix; UMatrix]) ; - add_unqualified ("trace_quad_form", ReturnType UReal, [UMatrix; UVector]) ; - add_unqualified ("trace_quad_form", ReturnType UReal, [UMatrix; UMatrix]) ; - add_unqualified ("transpose", ReturnType URowVector, [UVector]) ; - add_unqualified ("transpose", ReturnType UVector, [URowVector]) ; - add_unqualified ("transpose", ReturnType UMatrix, [UMatrix]) ; - add_unqualified ("uniform_simplex", ReturnType UVector, [UInt]) ; - add_unqualified ("variance", ReturnType UReal, [bare_array_type (UReal, 1)]) ; - add_unqualified ("variance", ReturnType UReal, [UVector]) ; - add_unqualified ("variance", ReturnType UReal, [URowVector]) ; - add_unqualified ("variance", ReturnType UReal, [UMatrix]) ; - add_unqualified ("wishart_rng", ReturnType UMatrix, [UReal; UMatrix]) ; - add_unqualified ("zeros_int_array", ReturnType (UArray UInt), [UInt]) ; - add_unqualified ("zeros_array", ReturnType (UArray UReal), [UInt]) ; - add_unqualified ("zeros_row_vector", ReturnType URowVector, [UInt]) ; - add_unqualified ("zeros_vector", ReturnType UVector, [UInt]) ; + , [ bare_array_type (UInt, i) ] ) + done; + add_unqualified ("to_array_2d", ReturnType (bare_array_type (UReal, 2)), [ UMatrix ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UMatrix; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UMatrix; UInt; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UVector ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UVector; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ UVector; UInt; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ URowVector ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ URowVector; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ URowVector; UInt; UInt; UInt ]); + add_unqualified + ("to_matrix", ReturnType UMatrix, [ bare_array_type (UReal, 1); UInt; UInt ]); + add_unqualified + ("to_matrix", ReturnType UMatrix, [ bare_array_type (UReal, 1); UInt; UInt; UInt ]); + add_unqualified + ("to_matrix", ReturnType UMatrix, [ bare_array_type (UInt, 1); UInt; UInt ]); + add_unqualified + ("to_matrix", ReturnType UMatrix, [ bare_array_type (UInt, 1); UInt; UInt; UInt ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ bare_array_type (UReal, 2) ]); + add_unqualified ("to_matrix", ReturnType UMatrix, [ bare_array_type (UInt, 2) ]); + add_unqualified ("to_row_vector", ReturnType URowVector, [ UMatrix ]); + add_unqualified ("to_row_vector", ReturnType URowVector, [ UVector ]); + add_unqualified ("to_row_vector", ReturnType URowVector, [ URowVector ]); + add_unqualified ("to_row_vector", ReturnType URowVector, [ bare_array_type (UReal, 1) ]); + add_unqualified ("to_row_vector", ReturnType URowVector, [ bare_array_type (UInt, 1) ]); + add_unqualified ("to_vector", ReturnType UVector, [ UMatrix ]); + add_unqualified ("to_vector", ReturnType UVector, [ UVector ]); + add_unqualified ("to_vector", ReturnType UVector, [ URowVector ]); + add_unqualified ("to_vector", ReturnType UVector, [ bare_array_type (UReal, 1) ]); + add_unqualified ("to_vector", ReturnType UVector, [ bare_array_type (UInt, 1) ]); + add_unqualified ("trace", ReturnType UReal, [ UMatrix ]); + add_unqualified ("trace_gen_quad_form", ReturnType UReal, [ UMatrix; UMatrix; UMatrix ]); + add_unqualified ("trace_quad_form", ReturnType UReal, [ UMatrix; UVector ]); + add_unqualified ("trace_quad_form", ReturnType UReal, [ UMatrix; UMatrix ]); + add_unqualified ("transpose", ReturnType URowVector, [ UVector ]); + add_unqualified ("transpose", ReturnType UVector, [ URowVector ]); + add_unqualified ("transpose", ReturnType UMatrix, [ UMatrix ]); + add_unqualified ("uniform_simplex", ReturnType UVector, [ UInt ]); + add_unqualified ("variance", ReturnType UReal, [ bare_array_type (UReal, 1) ]); + add_unqualified ("variance", ReturnType UReal, [ UVector ]); + add_unqualified ("variance", ReturnType UReal, [ URowVector ]); + add_unqualified ("variance", ReturnType UReal, [ UMatrix ]); + add_unqualified ("wishart_rng", ReturnType UMatrix, [ UReal; UMatrix ]); + add_unqualified ("zeros_int_array", ReturnType (UArray UInt), [ UInt ]); + add_unqualified ("zeros_array", ReturnType (UArray UReal), [ UInt ]); + add_unqualified ("zeros_row_vector", ReturnType URowVector, [ UInt ]); + add_unqualified ("zeros_vector", ReturnType UVector, [ UInt ]); (* Now add all the manually added stuff to the main hashtable used for type-checking *) Hashtbl.iteri manual_stan_math_signatures ~f:(fun ~key ~data -> - List.iter data ~f:(fun data -> - Hashtbl.add_multi stan_math_signatures ~key ~data ) ) + List.iter data ~f:(fun data -> Hashtbl.add_multi stan_math_signatures ~key ~data)) +;; diff --git a/src/middle/Stmt.ml b/src/middle/Stmt.ml index 2539198d6c..7e9c34aa10 100644 --- a/src/middle/Stmt.ml +++ b/src/middle/Stmt.ml @@ -17,13 +17,19 @@ module Fixed = struct | Skip | IfElse of 'a * 'b * 'b option | While of 'a * 'b - | For of {loopvar: string; lower: 'a; upper: 'a; body: 'b} + | For of + { loopvar : string + ; lower : 'a + ; upper : 'a + ; body : 'b + } | Block of 'b list | SList of 'b list | Decl of - { decl_adtype: UnsizedType.autodifftype - ; decl_id: string - ; decl_type: 'a Type.t } + { decl_adtype : UnsizedType.autodifftype + ; decl_id : string + ; decl_type : 'a Type.t + } [@@deriving sexp, hash, map, fold, compare] and 'a lvalue = string * UnsizedType.t * 'a Index.t list @@ -31,39 +37,59 @@ module Fixed = struct let pp pp_e pp_s ppf = function | Assignment ((assignee, _, idcs), rhs) -> - Fmt.pf ppf {|@[%a =@ %a;@]|} (Index.pp_indexed pp_e) - (assignee, idcs) pp_e rhs - | TargetPE expr -> - Fmt.pf ppf {|@[%a +=@ %a;@]|} pp_keyword "target" pp_e expr + Fmt.pf ppf {|@[%a =@ %a;@]|} (Index.pp_indexed pp_e) (assignee, idcs) pp_e rhs + | TargetPE expr -> Fmt.pf ppf {|@[%a +=@ %a;@]|} pp_keyword "target" pp_e expr | NRFunApp (_, name, args) -> - Fmt.pf ppf {|@[%s%a;@]|} name - Fmt.(list pp_e ~sep:comma |> parens) - args + Fmt.pf ppf {|@[%s%a;@]|} name Fmt.(list pp_e ~sep:comma |> parens) args | Break -> pp_keyword ppf "break;" | Continue -> pp_keyword ppf "continue;" | Skip -> pp_keyword ppf ";" - | Return (Some expr) -> - Fmt.pf ppf {|%a %a;|} pp_keyword "return" pp_e expr + | Return (Some expr) -> Fmt.pf ppf {|%a %a;|} pp_keyword "return" pp_e expr | Return _ -> pp_keyword ppf "return;" | IfElse (pred, s_true, Some s_false) -> - Fmt.pf ppf {|%a(%a) %a %a %a|} pp_builtin_syntax "if" pp_e pred pp_s - s_true pp_builtin_syntax "else" pp_s s_false + Fmt.pf + ppf + {|%a(%a) %a %a %a|} + pp_builtin_syntax + "if" + pp_e + pred + pp_s + s_true + pp_builtin_syntax + "else" + pp_s + s_false | IfElse (pred, s_true, _) -> - Fmt.pf ppf {|%a(%a) %a|} pp_builtin_syntax "if" pp_e pred pp_s s_true + Fmt.pf ppf {|%a(%a) %a|} pp_builtin_syntax "if" pp_e pred pp_s s_true | While (pred, stmt) -> - Fmt.pf ppf {|%a(%a) %a|} pp_builtin_syntax "while" pp_e pred pp_s - stmt - | For {loopvar; lower; upper; body} -> - Fmt.pf ppf {|%a(%s in %a:%a) %a|} pp_builtin_syntax "for" loopvar - pp_e lower pp_e upper pp_s body + Fmt.pf ppf {|%a(%a) %a|} pp_builtin_syntax "while" pp_e pred pp_s stmt + | For { loopvar; lower; upper; body } -> + Fmt.pf + ppf + {|%a(%s in %a:%a) %a|} + pp_builtin_syntax + "for" + loopvar + pp_e + lower + pp_e + upper + pp_s + body | Block stmts -> - Fmt.pf ppf {|{@;<1 2>@[%a@]@;}|} - Fmt.(list pp_s ~sep:Fmt.cut) - stmts + Fmt.pf ppf {|{@;<1 2>@[%a@]@;}|} Fmt.(list pp_s ~sep:Fmt.cut) stmts | SList stmts -> Fmt.(list pp_s ~sep:Fmt.cut |> vbox) ppf stmts - | Decl {decl_adtype; decl_id; decl_type} -> - Fmt.pf ppf {|%a%a %s;|} UnsizedType.pp_autodifftype decl_adtype - (Type.pp pp_e) decl_type decl_id + | Decl { decl_adtype; decl_id; decl_type } -> + Fmt.pf + ppf + {|%a%a %s;|} + UnsizedType.pp_autodifftype + decl_adtype + (Type.pp pp_e) + decl_type + decl_id + ;; include Foldable.Make2 (struct type nonrec ('a, 'b) t = ('a, 'b) t @@ -101,7 +127,7 @@ module Located = struct include Specialized.Make2 (Fixed) (Expr.Typed) (Meta) - let loc_of Fixed.({meta; _}) = meta + let loc_of Fixed.{ meta; _ } = meta (** This module acts as a temporary replace for the `stmt_loc_num` type that is currently used within `analysis_and_optimization`. @@ -114,8 +140,9 @@ module Located = struct *) module Non_recursive = struct type t = - { pattern: (Expr.Typed.t, int) Fixed.Pattern.t - ; meta: Meta.t sexp_opaque [@compare.ignore] } + { pattern : (Expr.Typed.t, int) Fixed.Pattern.t + ; meta : Meta.t sexp_opaque [@compare.ignore] + } [@@deriving compare, sexp, hash] end end @@ -125,93 +152,90 @@ both are typed and labelled. *) module Labelled = struct module Meta = struct type t = - { loc: Location_span.t sexp_opaque [@compare.ignore] - ; label: Label.Int_label.t [@compare.ignore] } + { loc : Location_span.t sexp_opaque [@compare.ignore] + ; label : Label.Int_label.t [@compare.ignore] + } [@@deriving compare, create, sexp, hash] - let empty = - create ~loc:Location_span.empty ~label:Label.Int_label.(prev init) () - + let empty = create ~loc:Location_span.empty ~label:Label.Int_label.(prev init) () let pp _ _ = () end include Specialized.Make2 (Fixed) (Expr.Labelled) (Meta) - let label_of Fixed.({meta= Meta.({label; _}); _}) = label - let loc_of Fixed.({meta= Meta.({loc; _}); _}) = loc + let label_of Fixed.{ meta = Meta.{ label; _ }; _ } = label + let loc_of Fixed.{ meta = Meta.{ loc; _ }; _ } = loc let label ?(init = Label.Int_label.init) (stmt : Located.t) : t = let lbl = ref init in - let f Expr.Typed.Meta.({adlevel; type_; loc}) = + let f Expr.Typed.Meta.{ adlevel; type_; loc } = let cur_lbl = !lbl in - lbl := Label.Int_label.next cur_lbl ; + lbl := Label.Int_label.next cur_lbl; Expr.Labelled.Meta.create ~type_ ~loc ~adlevel ~label:cur_lbl () and g loc = let cur_lbl = !lbl in - lbl := Label.Int_label.next cur_lbl ; + lbl := Label.Int_label.next cur_lbl; Meta.create ~loc ~label:cur_lbl () in Fixed.map f g stmt + ;; type associations = - { exprs: Expr.Labelled.t Label.Int_label.Map.t - ; stmts: t Label.Int_label.Map.t } + { exprs : Expr.Labelled.t Label.Int_label.Map.t + ; stmts : t Label.Int_label.Map.t + } - let empty = - {exprs= Label.Int_label.Map.empty; stmts= Label.Int_label.Map.empty} + let empty = { exprs = Label.Int_label.Map.empty; stmts = Label.Int_label.Map.empty } - let rec associate ?init:(assocs = empty) ({pattern; _} as stmt : t) = + let rec associate ?init:(assocs = empty) ({ pattern; _ } as stmt : t) = associate_pattern { assocs with - stmts= - Label.Int_label.Map.add_exn assocs.stmts ~key:(label_of stmt) - ~data:stmt } + stmts = Label.Int_label.Map.add_exn assocs.stmts ~key:(label_of stmt) ~data:stmt + } pattern and associate_pattern assocs = function | Fixed.Pattern.Break | Skip | Continue | Return None -> assocs | Return (Some e) | TargetPE e -> - {assocs with exprs= Expr.Labelled.associate ~init:assocs.exprs e} + { assocs with exprs = Expr.Labelled.associate ~init:assocs.exprs e } | NRFunApp (_, _, args) -> - { assocs with - exprs= - List.fold args ~init:assocs.exprs ~f:(fun accu x -> - Expr.Labelled.associate ~init:accu x ) } + { assocs with + exprs = + List.fold args ~init:assocs.exprs ~f:(fun accu x -> + Expr.Labelled.associate ~init:accu x) + } | Assignment ((_, _, idxs), rhs) -> - let exprs = - Expr.Labelled.( - associate rhs - ~init:(List.fold ~f:associate_index ~init:assocs.exprs idxs)) - in - {assocs with exprs} + let exprs = + Expr.Labelled.( + associate rhs ~init:(List.fold ~f:associate_index ~init:assocs.exprs idxs)) + in + { assocs with exprs } | IfElse (pred, body, None) | While (pred, body) -> - let exprs = Expr.Labelled.associate ~init:assocs.exprs pred in - associate ~init:{assocs with exprs} body + let exprs = Expr.Labelled.associate ~init:assocs.exprs pred in + associate ~init:{ assocs with exprs } body | IfElse (pred, ts, Some fs) -> - let exprs = Expr.Labelled.associate ~init:assocs.exprs pred in - let assocs' = {assocs with exprs} in - associate ~init:(associate ~init:assocs' ts) fs - | Decl {decl_type; _} -> associate_possibly_sized_type assocs decl_type - | For {lower; upper; body; _} -> - let exprs = - Expr.Labelled.( - associate ~init:(associate ~init:assocs.exprs lower) upper) - in - let assocs' = {assocs with exprs} in - associate ~init:assocs' body + let exprs = Expr.Labelled.associate ~init:assocs.exprs pred in + let assocs' = { assocs with exprs } in + associate ~init:(associate ~init:assocs' ts) fs + | Decl { decl_type; _ } -> associate_possibly_sized_type assocs decl_type + | For { lower; upper; body; _ } -> + let exprs = + Expr.Labelled.(associate ~init:(associate ~init:assocs.exprs lower) upper) + in + let assocs' = { assocs with exprs } in + associate ~init:assocs' body | Block xs | SList xs -> - List.fold ~f:(fun accu x -> associate ~init:accu x) ~init:assocs xs + List.fold ~f:(fun accu x -> associate ~init:accu x) ~init:assocs xs and associate_possibly_sized_type assocs = function - | Type.Sized st -> - {assocs with exprs= SizedType.associate ~init:assocs.exprs st} + | Type.Sized st -> { assocs with exprs = SizedType.associate ~init:assocs.exprs st } | Unsized _ -> assocs + ;; end module Numbered = struct module Meta = struct - type t = (int sexp_opaque[@compare.ignore]) - [@@deriving compare, sexp, hash] + type t = (int sexp_opaque[@compare.ignore]) [@@deriving compare, sexp, hash] let empty = 0 let from_int (i : int) : t = i @@ -224,84 +248,87 @@ end module Helpers = struct let ensure_var bodyfn (expr : Expr.Typed.t) meta = match expr with - | {pattern= Var _; _} -> bodyfn expr meta + | { pattern = Var _; _ } -> bodyfn expr meta | _ -> - let symbol, reset = Gensym.enter () in - let body = bodyfn {expr with pattern= Var symbol} meta in - let decl = - { body with - Fixed.pattern= - Decl - { decl_adtype= Expr.Typed.adlevel_of expr - ; decl_id= symbol - ; decl_type= Unsized (Expr.Typed.type_of expr) } } - in - let assign = - { body with - Fixed.pattern= - Assignment ((symbol, Expr.Typed.type_of expr, []), expr) } - in - reset () ; - {body with Fixed.pattern= Block [decl; assign; body]} + let symbol, reset = Gensym.enter () in + let body = bodyfn { expr with pattern = Var symbol } meta in + let decl = + { body with + Fixed.pattern = + Decl + { decl_adtype = Expr.Typed.adlevel_of expr + ; decl_id = symbol + ; decl_type = Unsized (Expr.Typed.type_of expr) + } + } + in + let assign = + { body with + Fixed.pattern = Assignment ((symbol, Expr.Typed.type_of expr, []), expr) + } + in + reset (); + { body with Fixed.pattern = Block [ decl; assign; body ] } + ;; let internal_nrfunapp fn args meta = - { Fixed.pattern= - NRFunApp (CompilerInternal, Internal_fun.to_string fn, args) - ; meta } + { Fixed.pattern = NRFunApp (CompilerInternal, Internal_fun.to_string fn, args); meta } + ;; (** [mkfor] returns a MIR For statement that iterates over the given expression [iteratee]. *) let mkfor upper bodyfn iteratee meta = let idx s = - let meta = - Expr.Typed.Meta.create ~type_:UInt ~loc:meta ~adlevel:DataOnly () - in - let expr = Expr.Fixed.{meta; pattern= Var s} in + let meta = Expr.Typed.Meta.create ~type_:UInt ~loc:meta ~adlevel:DataOnly () in + let expr = Expr.Fixed.{ meta; pattern = Var s } in Index.Single expr in let loopvar, reset = Gensym.enter () in let lower = Expr.Helpers.loop_bottom in let stmt = - Fixed.Pattern.Block - [bodyfn (Expr.Helpers.add_int_index iteratee (idx loopvar))] + Fixed.Pattern.Block [ bodyfn (Expr.Helpers.add_int_index iteratee (idx loopvar)) ] in - reset () ; - let body = Fixed.{meta; pattern= stmt} in - let pattern = Fixed.Pattern.For {loopvar; lower; upper; body} in - Fixed.{meta; pattern} + reset (); + let body = Fixed.{ meta; pattern = stmt } in + let pattern = Fixed.Pattern.For { loopvar; lower; upper; body } in + Fixed.{ meta; pattern } + ;; let rec for_each bodyfn iteratee smeta = let len (e : Expr.Typed.t) = let emeta = e.meta in - let emeta' = {emeta with Expr.Typed.Meta.type_= UInt} in - Expr.Helpers.internal_funapp FnLength [e] emeta' + let emeta' = { emeta with Expr.Typed.Meta.type_ = UInt } in + Expr.Helpers.internal_funapp FnLength [ e ] emeta' in match Expr.Typed.type_of iteratee with | UInt | UReal -> bodyfn iteratee | UVector | URowVector -> mkfor (len iteratee) bodyfn iteratee smeta | UMatrix -> - let emeta = iteratee.meta in - let emeta' = {emeta with Expr.Typed.Meta.type_= UInt} in - let rows = - Expr.Fixed. - {meta= emeta'; pattern= FunApp (StanLib, "rows", [iteratee])} - in - mkfor rows (fun e -> for_each bodyfn e smeta) iteratee smeta + let emeta = iteratee.meta in + let emeta' = { emeta with Expr.Typed.Meta.type_ = UInt } in + let rows = + Expr.Fixed.{ meta = emeta'; pattern = FunApp (StanLib, "rows", [ iteratee ]) } + in + mkfor rows (fun e -> for_each bodyfn e smeta) iteratee smeta | UArray _ -> mkfor (len iteratee) bodyfn iteratee smeta | UMathLibraryFunction | UFun _ -> - raise_s [%message "can't iterate over " (iteratee : Expr.Typed.t)] + raise_s [%message "can't iterate over " (iteratee : Expr.Typed.t)] + ;; let contains_fn fn ?(init = false) stmt = let fstr = Internal_fun.to_string fn in - let rec aux accu Fixed.({pattern; _}) = + let rec aux accu Fixed.{ pattern; _ } = match pattern with | NRFunApp (_, fname, _) when fname = fstr -> true | stmt_pattern -> - Fixed.Pattern.fold_left ~init:accu stmt_pattern - ~f:(fun accu expr -> Expr.Helpers.contains_fn fn ~init:accu expr) - ~g:aux + Fixed.Pattern.fold_left + ~init:accu + stmt_pattern + ~f:(fun accu expr -> Expr.Helpers.contains_fn fn ~init:accu expr) + ~g:aux in aux init stmt + ;; (** [for_eigen unsizedtype...] generates a For statement that loops over the eigen types in the underlying [unsizedtype]; i.e. just iterating @@ -315,9 +342,9 @@ module Helpers = struct *) let rec for_eigen st bodyfn var smeta = match st with - | SizedType.SInt | SReal | SVector _ | SRowVector _ | SMatrix _ -> - bodyfn var + | SizedType.SInt | SReal | SVector _ | SRowVector _ | SMatrix _ -> bodyfn var | SArray (t, d) -> mkfor d (fun e -> for_eigen t bodyfn e smeta) var smeta + ;; (** [for_scalar unsizedtype...] generates a For statement that loops over the scalars in the underlying [unsizedtype]. @@ -333,32 +360,35 @@ module Helpers = struct | SizedType.SInt | SReal -> bodyfn var | SVector d | SRowVector d -> mkfor d bodyfn var smeta | SMatrix (d1, d2) -> - mkfor d1 (fun e -> for_scalar (SRowVector d2) bodyfn e smeta) var smeta + mkfor d1 (fun e -> for_scalar (SRowVector d2) bodyfn e smeta) var smeta | SArray (t, d) -> mkfor d (fun e -> for_scalar t bodyfn e smeta) var smeta + ;; (** Exactly like for_scalar, but iterating through array dimensions in the inverted order.*) let for_scalar_inv st bodyfn (var : Expr.Typed.t) smeta = - let var = {var with pattern= Indexed (var, [])} in - let invert_index_order (Expr.Fixed.({pattern; _}) as e) = + let var = { var with pattern = Indexed (var, []) } in + let invert_index_order (Expr.Fixed.{ pattern; _ } as e) = match pattern with | Indexed (obj, []) -> obj - | Indexed (obj, idxs) -> {e with pattern= Indexed (obj, List.rev idxs)} + | Indexed (obj, idxs) -> { e with pattern = Indexed (obj, List.rev idxs) } | _ -> e in let rec go st bodyfn var smeta = match st with | SizedType.SArray (t, d) -> - let bodyfn' var = mkfor d bodyfn var smeta in - go t bodyfn' var smeta + let bodyfn' var = mkfor d bodyfn var smeta in + go t bodyfn' var smeta | SMatrix (d1, d2) -> - let bodyfn' var = mkfor d1 bodyfn var smeta in - go (SRowVector d2) bodyfn' var smeta + let bodyfn' var = mkfor d1 bodyfn var smeta in + go (SRowVector d2) bodyfn' var smeta | _ -> for_scalar st bodyfn var smeta in go st (Fn.compose bodyfn invert_index_order) var smeta + ;; let assign_indexed decl_type vident meta varfn var = let indices = Expr.Helpers.collect_indices var in - Fixed.{meta; pattern= Assignment ((vident, decl_type, indices), varfn var)} + Fixed.{ meta; pattern = Assignment ((vident, decl_type, indices), varfn var) } + ;; end diff --git a/src/middle/Stmt.mli b/src/middle/Stmt.mli index f6a2527d0c..de8069f0e1 100644 --- a/src/middle/Stmt.mli +++ b/src/middle/Stmt.mli @@ -14,13 +14,19 @@ module Fixed : sig | Skip | IfElse of 'a * 'b * 'b option | While of 'a * 'b - | For of {loopvar: string; lower: 'a; upper: 'a; body: 'b} + | For of + { loopvar : string + ; lower : 'a + ; upper : 'a + ; body : 'b + } | Block of 'b list | SList of 'b list | Decl of - { decl_adtype: UnsizedType.autodifftype - ; decl_id: string - ; decl_type: 'a Type.t } + { decl_adtype : UnsizedType.autodifftype + ; decl_id : string + ; decl_type : 'a Type.t + } [@@deriving sexp, hash, compare] and 'a lvalue = string * UnsizedType.t * 'a Index.t list @@ -41,8 +47,8 @@ module NoMeta : sig include Specialized.S - with module Meta := Meta - and type t = (Expr.NoMeta.Meta.t, Meta.t) Fixed.t + with module Meta := Meta + and type t = (Expr.NoMeta.Meta.t, Meta.t) Fixed.t val remove_meta : ('a, 'b) Fixed.t -> t end @@ -57,18 +63,16 @@ module Located : sig include Specialized.S - with module Meta := Meta - and type t = - ( Expr.Typed.Meta.t - , (Meta.t sexp_opaque[@compare.ignore]) ) - Fixed.t + with module Meta := Meta + and type t = (Expr.Typed.Meta.t, (Meta.t sexp_opaque[@compare.ignore])) Fixed.t val loc_of : t -> Location_span.t module Non_recursive : sig type t = - { pattern: (Expr.Typed.t, int) Fixed.Pattern.t - ; meta: Meta.t sexp_opaque [@compare.ignore] } + { pattern : (Expr.Typed.t, int) Fixed.Pattern.t + ; meta : Meta.t sexp_opaque [@compare.ignore] + } [@@deriving compare, sexp, hash] end end @@ -76,8 +80,9 @@ end module Labelled : sig module Meta : sig type t = - { loc: Location_span.t sexp_opaque [@compare.ignore] - ; label: Int_label.t [@compare.ignore] } + { loc : Location_span.t sexp_opaque [@compare.ignore] + ; label : Int_label.t [@compare.ignore] + } [@@deriving compare, create, sexp, hash] include Specialized.Meta with type t := t @@ -85,23 +90,24 @@ module Labelled : sig include Specialized.S - with module Meta := Meta - and type t = (Expr.Labelled.Meta.t, Meta.t) Fixed.t + with module Meta := Meta + and type t = (Expr.Labelled.Meta.t, Meta.t) Fixed.t val loc_of : t -> Location_span.t val label_of : t -> Int_label.t val label : ?init:int -> Located.t -> t type associations = - {exprs: Expr.Labelled.t Int_label.Map.t; stmts: t Int_label.Map.t} + { exprs : Expr.Labelled.t Int_label.Map.t + ; stmts : t Int_label.Map.t + } val associate : ?init:associations -> t -> associations end module Numbered : sig module Meta : sig - type t = (int sexp_opaque[@compare.ignore]) - [@@deriving compare, sexp, hash] + type t = (int sexp_opaque[@compare.ignore]) [@@deriving compare, sexp, hash] include Specialized.Meta with type t := t @@ -110,52 +116,57 @@ module Numbered : sig include Specialized.S - with module Meta := Meta - and type t = (Expr.Typed.Meta.t, Meta.t) Fixed.t + with module Meta := Meta + and type t = (Expr.Typed.Meta.t, Meta.t) Fixed.t end module Helpers : sig - val ensure_var : - (Expr.Typed.t -> 'a -> Located.t) -> Expr.Typed.t -> 'a -> Located.t + val ensure_var : (Expr.Typed.t -> 'a -> Located.t) -> Expr.Typed.t -> 'a -> Located.t - val internal_nrfunapp : - Internal_fun.t -> 'a Fixed.First.t list -> 'b -> ('a, 'b) Fixed.t + val internal_nrfunapp + : Internal_fun.t + -> 'a Fixed.First.t list + -> 'b + -> ('a, 'b) Fixed.t val contains_fn : Internal_fun.t -> ?init:bool -> ('a, 'b) Fixed.t -> bool - val mkfor : - Expr.Typed.t + val mkfor + : Expr.Typed.t -> (Expr.Typed.t -> Located.t) -> Expr.Typed.t -> Location_span.t -> Located.t - val for_each : - (Expr.Typed.t -> Located.t) -> Expr.Typed.t -> Location_span.t -> Located.t + val for_each + : (Expr.Typed.t -> Located.t) + -> Expr.Typed.t + -> Location_span.t + -> Located.t - val for_scalar : - Expr.Typed.t SizedType.t + val for_scalar + : Expr.Typed.t SizedType.t -> (Expr.Typed.t -> Located.t) -> Expr.Typed.t -> Location_span.t -> Located.t - val for_scalar_inv : - Expr.Typed.t SizedType.t + val for_scalar_inv + : Expr.Typed.t SizedType.t -> (Expr.Typed.t -> Located.t) -> Expr.Typed.t -> Location_span.t -> Located.t - val for_eigen : - Expr.Typed.t SizedType.t + val for_eigen + : Expr.Typed.t SizedType.t -> (Expr.Typed.t -> Located.t) -> Expr.Typed.t -> Location_span.t -> Located.t - val assign_indexed : - UnsizedType.t + val assign_indexed + : UnsizedType.t -> string -> 'a -> ('b Expr.Fixed.t -> 'b Expr.Fixed.t) diff --git a/src/middle/Type.ml b/src/middle/Type.ml index 04eb3e497e..e33e11043a 100644 --- a/src/middle/Type.ml +++ b/src/middle/Type.ml @@ -1,18 +1,26 @@ open Common -type 'a t = Sized of 'a SizedType.t | Unsized of UnsizedType.t +type 'a t = + | Sized of 'a SizedType.t + | Unsized of UnsizedType.t [@@deriving sexp, compare, map, hash, fold] let pp pp_e ppf = function | Sized st -> SizedType.pp pp_e ppf st | Unsized ust -> UnsizedType.pp ppf ust +;; -let collect_exprs = function Sized st -> SizedType.collect_exprs st | _ -> [] +let collect_exprs = function + | Sized st -> SizedType.collect_exprs st + | _ -> [] +;; let to_unsized = function | Sized st -> SizedType.to_unsized st | Unsized ut -> ut +;; let associate ?init:(assocs = Label.Int_label.Map.empty) = function | Sized st -> SizedType.associate ~init:assocs st | Unsized _ -> assocs +;; diff --git a/src/middle/UnsizedType.ml b/src/middle/UnsizedType.ml index fae0b7284f..83433c8011 100644 --- a/src/middle/UnsizedType.ml +++ b/src/middle/UnsizedType.ml @@ -11,20 +11,27 @@ type t = | UFun of (autodifftype * t) list * returntype | UMathLibraryFunction -and autodifftype = DataOnly | AutoDiffable +and autodifftype = + | DataOnly + | AutoDiffable -and returntype = Void | ReturnType of t [@@deriving compare, hash, sexp] +and returntype = + | Void + | ReturnType of t +[@@deriving compare, hash, sexp] let pp_autodifftype ppf = function | DataOnly -> pp_keyword ppf "data " | AutoDiffable -> () +;; let unsized_array_depth unsized_ty = let rec aux depth = function | UArray ut -> aux (depth + 1) ut - | ut -> (ut, depth) + | ut -> ut, depth in aux 0 unsized_ty +;; let rec pp ppf = function | UInt -> pp_keyword ppf "int" @@ -33,15 +40,18 @@ let rec pp ppf = function | URowVector -> pp_keyword ppf "row_vector" | UMatrix -> pp_keyword ppf "matrix" | UArray ut -> - let ty, depth = unsized_array_depth ut in - let commas = String.make depth ',' in - Fmt.pf ppf "%a[%s]" pp ty commas + let ty, depth = unsized_array_depth ut in + let commas = String.make depth ',' in + Fmt.pf ppf "%a[%s]" pp ty commas | UFun (argtypes, rt) -> - Fmt.pf ppf {|@[(%a) => %a@]|} - Fmt.(list pp_fun_arg ~sep:comma) - argtypes pp_returntype rt - | UMathLibraryFunction -> - (pp_angle_brackets Fmt.string) ppf "Stan Math function" + Fmt.pf + ppf + {|@[(%a) => %a@]|} + Fmt.(list pp_fun_arg ~sep:comma) + argtypes + pp_returntype + rt + | UMathLibraryFunction -> (pp_angle_brackets Fmt.string) ppf "Stan Math function" and pp_fun_arg ppf (ad_ty, unsized_ty) = match ad_ty with @@ -51,37 +61,45 @@ and pp_fun_arg ppf (ad_ty, unsized_ty) = and pp_returntype ppf = function | Void -> Fmt.string ppf "void" | ReturnType ut -> pp ppf ut +;; (* -- Type conversion -- *) let autodifftype_can_convert at1 at2 = - match (at1, at2) with DataOnly, AutoDiffable -> false | _ -> true + match at1, at2 with + | DataOnly, AutoDiffable -> false + | _ -> true +;; let rec lub_ad_type = function | [] -> DataOnly | x :: xs -> - let y = lub_ad_type xs in - if compare_autodifftype x y < 0 then y else x + let y = lub_ad_type xs in + if compare_autodifftype x y < 0 then y else x +;; let check_of_same_type_mod_conv name t1 t2 = - if String.is_prefix name ~prefix:"assign_" then t1 = t2 - else - match (t1, t2) with + if String.is_prefix name ~prefix:"assign_" + then t1 = t2 + else ( + match t1, t2 with | UReal, UInt -> true | UFun (l1, rt1), UFun (l2, rt2) -> - rt1 = rt2 - && List.for_all - ~f:(fun x -> x = true) - (List.map2_exn - ~f:(fun (at1, ut1) (at2, ut2) -> - ut1 = ut2 && autodifftype_can_convert at2 at1 ) - l1 l2) - | _ -> t1 = t2 + rt1 = rt2 + && List.for_all + ~f:(fun x -> x = true) + (List.map2_exn + ~f:(fun (at1, ut1) (at2, ut2) -> + ut1 = ut2 && autodifftype_can_convert at2 at1) + l1 + l2) + | _ -> t1 = t2) +;; let rec check_of_same_type_mod_array_conv name t1 t2 = - match (t1, t2) with - | UArray t1elt, UArray t2elt -> - check_of_same_type_mod_array_conv name t1elt t2elt + match t1, t2 with + | UArray t1elt, UArray t2elt -> check_of_same_type_mod_array_conv name t1elt t2elt | _ -> check_of_same_type_mod_conv name t1 t2 +;; let check_compatible_arguments_mod_conv name args1 args2 = List.length args1 = List.length args2 @@ -90,45 +108,64 @@ let check_compatible_arguments_mod_conv name args1 args2 = (List.map2_exn ~f:(fun sign1 sign2 -> check_of_same_type_mod_conv name (snd sign1) (snd sign2) - && autodifftype_can_convert (fst sign1) (fst sign2) ) - args1 args2) + && autodifftype_can_convert (fst sign1) (fst sign2)) + args1 + args2) +;; (** Given two types find the minimal type both can convert to *) let rec common_type = function | UReal, UInt | UInt, UReal -> Some UReal - | UArray t1, UArray t2 -> - common_type (t1, t2) |> Option.map ~f:(fun t -> UArray t) + | UArray t1, UArray t2 -> common_type (t1, t2) |> Option.map ~f:(fun t -> UArray t) | t1, t2 when t1 = t2 -> Some t1 | _, _ -> None +;; (* -- Helpers -- *) let is_real_type = function | UReal | UVector | URowVector | UMatrix - |UArray UReal - |UArray UVector - |UArray URowVector - |UArray UMatrix -> - true + | UArray UReal + | UArray UVector + | UArray URowVector + | UArray UMatrix -> true | _ -> false +;; let rec is_autodiffable = function | UReal | UVector | URowVector | UMatrix -> true | UArray t -> is_autodiffable t | _ -> false +;; -let is_scalar_type = function UReal | UInt -> true | _ -> false -let is_int_type = function UInt | UArray UInt -> true | _ -> false -let is_fun_type = function UFun _ -> true | _ -> false +let is_scalar_type = function + | UReal | UInt -> true + | _ -> false +;; + +let is_int_type = function + | UInt | UArray UInt -> true + | _ -> false +;; + +let is_fun_type = function + | UFun _ -> true + | _ -> false +;; (** Detect if type contains an integer *) let rec contains_int ut = - match ut with UInt -> true | UArray ut -> contains_int ut | _ -> false + 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 | UMatrix, _ -> true | _ -> false +;; module Comparator = Comparator.Make (struct type nonrec t = t diff --git a/src/middle/Utils.ml b/src/middle/Utils.ml index c95304eb5d..547fb88a4b 100644 --- a/src/middle/Utils.ml +++ b/src/middle/Utils.ml @@ -4,48 +4,52 @@ let option_or_else ~if_none x = Option.first_some x if_none (* Name mangling helper functions for distributions *) let proportional_to_distribution_infix = "_propto" -let distribution_suffices = ["_lpmf"; "_lpdf"; "_log"] +let distribution_suffices = [ "_lpmf"; "_lpdf"; "_log" ] let propto_suffices = - List.map - ~f:(fun x -> proportional_to_distribution_infix ^ x) - distribution_suffices + List.map ~f:(fun x -> proportional_to_distribution_infix ^ x) distribution_suffices +;; let is_user_ident = Fn.non (String.is_suffix ~suffix:"__") let is_distribution_name s = - (not - ( String.is_suffix s ~suffix:"_cdf_log" - || String.is_suffix s ~suffix:"_ccdf_log" )) - && List.exists - ~f:(fun suffix -> String.is_suffix s ~suffix) - distribution_suffices + (not (String.is_suffix s ~suffix:"_cdf_log" || String.is_suffix s ~suffix:"_ccdf_log")) + && List.exists ~f:(fun suffix -> String.is_suffix s ~suffix) distribution_suffices +;; let is_propto_distribution s = List.exists ~f:(fun suffix -> String.is_suffix s ~suffix) propto_suffices +;; let remove_propto_infix suffix ~name = name |> String.chop_suffix ~suffix:(proportional_to_distribution_infix ^ suffix) |> Option.map ~f:(fun x -> x ^ suffix) +;; let stdlib_distribution_name s = List.map ~f:(remove_propto_infix ~name:s) distribution_suffices - |> List.filter_opt |> List.hd |> Option.value ~default:s + |> List.filter_opt + |> List.hd + |> Option.value ~default:s +;; let%expect_test "propto name mangling" = - stdlib_distribution_name "bernoulli_logit_propto_lpmf" |> print_string ; - stdlib_distribution_name "normal_propto_lpdf" |> ( ^ ) "; " |> print_string ; - stdlib_distribution_name "normal_lpdf" |> ( ^ ) "; " |> print_string ; - stdlib_distribution_name "normal" |> ( ^ ) "; " |> print_string ; + stdlib_distribution_name "bernoulli_logit_propto_lpmf" |> print_string; + stdlib_distribution_name "normal_propto_lpdf" |> ( ^ ) "; " |> print_string; + stdlib_distribution_name "normal_lpdf" |> ( ^ ) "; " |> print_string; + stdlib_distribution_name "normal" |> ( ^ ) "; " |> print_string; [%expect {| bernoulli_logit_lpmf; normal_lpdf; normal_lpdf; normal |}] +;; let all_but_last_n l n = List.fold_right l ~init:([], n) ~f:(fun ele (accum, n) -> - if n = 0 then (ele :: accum, n) else (accum, n - 1) ) + if n = 0 then ele :: accum, n else accum, n - 1) |> fst +;; let%expect_test "all but last n" = - let l = all_but_last_n [1; 2; 3; 4] 2 in - print_s [%sexp (l : int list)] ; + let l = all_but_last_n [ 1; 2; 3; 4 ] 2 in + print_s [%sexp (l : int list)]; [%expect {| (1 2) |}] +;; diff --git a/src/stan2tfp/Stan2tfp.ml b/src/stan2tfp/Stan2tfp.ml index d8afc0f8ea..ccaec328dc 100644 --- a/src/stan2tfp/Stan2tfp.ml +++ b/src/stan2tfp/Stan2tfp.ml @@ -12,7 +12,9 @@ let options = , "Dump the MIR after it's been transformed by the TFP backend." ) ; ( "--dump-mir" , Arg.Set dump_mir - , "Dump the MIR immediately after transformation from the AST." ) ] + , "Dump the MIR immediately after transformation from the AST." ) + ] +;; let usage = "Usage: stan2tfp [option] ... " let model_file = ref "" @@ -21,24 +23,25 @@ let remove_dotstan s = String.drop_suffix s 5 let set_model_file s = match !model_file with | "" -> - model_file := s ; - Semantic_check.model_name := - remove_dotstan (Filename.basename s) ^ "_model" + model_file := s; + Semantic_check.model_name := remove_dotstan (Filename.basename s) ^ "_model" | _ -> raise_s [%message "Can only pass in one model file."] +;; let main () = - Arg.parse options set_model_file usage ; + Arg.parse options set_model_file usage; let mir = - !model_file |> Frontend_utils.get_ast_or_exit + !model_file + |> Frontend_utils.get_ast_or_exit |> Frontend_utils.type_ast_or_exit |> Ast_to_Mir.trans_prog !Semantic_check.model_name in - if !dump_mir then - mir |> Middle.Program.Typed.sexp_of_t |> Sexp.to_string_hum - |> print_endline ; + if !dump_mir + then mir |> Middle.Program.Typed.sexp_of_t |> Sexp.to_string_hum |> print_endline; let mir = Transform_mir.trans_prog mir in - if !dump_transformed_mir then Fmt.pr "%a" Middle.Program.Typed.pp mir ; + if !dump_transformed_mir then Fmt.pr "%a" Middle.Program.Typed.pp mir; Fmt.pr "%a" Code_gen.pp_prog mir +;; let () = main () diff --git a/src/stan_math_backend/Cpp_Json.ml b/src/stan_math_backend/Cpp_Json.ml index 3899c914f5..120fcd6e9f 100644 --- a/src/stan_math_backend/Cpp_Json.ml +++ b/src/stan_math_backend/Cpp_Json.ml @@ -8,36 +8,43 @@ let rec sizedtype_to_json (st : Expr.Typed.t SizedType.t) : Yojson.Basic.t = |> Str.global_replace (Str.regexp "[\n\r\t ]+") " " in match st with - | SInt -> `Assoc [("name", `String "int")] - | SReal -> `Assoc [("name", `String "real")] + | SInt -> `Assoc [ "name", `String "int" ] + | SReal -> `Assoc [ "name", `String "real" ] | SVector d | SRowVector d -> - `Assoc [("name", `String "vector"); ("length", `String (emit_cpp_expr d))] + `Assoc [ "name", `String "vector"; "length", `String (emit_cpp_expr d) ] | SMatrix (d1, d2) -> - `Assoc - [ ("name", `String "matrix") - ; ("rows", `String (emit_cpp_expr d1)) - ; ("cols", `String (emit_cpp_expr d2)) ] + `Assoc + [ "name", `String "matrix" + ; "rows", `String (emit_cpp_expr d1) + ; "cols", `String (emit_cpp_expr d2) + ] | SArray (st, d) -> - `Assoc - [ ("name", `String "array") - ; ("length", `String (emit_cpp_expr d)) - ; ("element_type", sizedtype_to_json st) ] + `Assoc + [ "name", `String "array" + ; "length", `String (emit_cpp_expr d) + ; "element_type", sizedtype_to_json st + ] +;; let out_var_json (name, st, block) : Yojson.Basic.t = `Assoc - [ ("name", `String name) - ; ("type", sizedtype_to_json st) - ; ("block", `String (Fmt.strf "%a" Program.pp_io_block block)) ] + [ "name", `String name + ; "type", sizedtype_to_json st + ; "block", `String (Fmt.strf "%a" Program.pp_io_block block) + ] +;; let%expect_test "outvar to json pretty" = - let var x = {Expr.Fixed.pattern= Var x; meta= Expr.Typed.Meta.empty} in + let var x = { Expr.Fixed.pattern = Var x; meta = Expr.Typed.Meta.empty } in (* the following is equivalent to: parameters { vector[N] var_one[K]; } *) ("var_one", SArray (SVector (var "N"), var "K"), Parameters) - |> out_var_json |> Yojson.Basic.pretty_to_string |> print_endline ; + |> out_var_json + |> Yojson.Basic.pretty_to_string + |> print_endline; [%expect {| { @@ -49,25 +56,30 @@ let%expect_test "outvar to json pretty" = }, "block": "parameters" } |}] +;; let replace_cpp_expr s = s |> Str.global_replace (Str.regexp {|"|}) {|\"|} |> Str.global_replace (Str.regexp {|\\"<<|}) {|" <<|} |> Str.global_replace (Str.regexp {|>>\\"|}) {|<< "|} +;; let wrap_in_quotes s = "\"" ^ s ^ "\"" let out_var_interpolated_json_str vars = `List (List.map ~f:out_var_json vars) - |> Yojson.Basic.to_string |> replace_cpp_expr |> wrap_in_quotes + |> Yojson.Basic.to_string + |> replace_cpp_expr + |> wrap_in_quotes +;; let%expect_test "outvar to json" = - let var x = {Expr.Fixed.pattern= Var x; meta= Expr.Typed.Meta.empty} in - [ ( "var_one" - , SizedType.SArray (SVector (var "N"), var "K") - , Program.Parameters ) ] - |> out_var_interpolated_json_str |> print_endline ; + let var x = { Expr.Fixed.pattern = Var x; meta = Expr.Typed.Meta.empty } in + [ "var_one", SizedType.SArray (SVector (var "N"), var "K"), Program.Parameters ] + |> out_var_interpolated_json_str + |> print_endline; [%expect {| "[{\"name\":\"var_one\",\"type\":{\"name\":\"array\",\"length\":" << K << ",\"element_type\":{\"name\":\"vector\",\"length\":" << N << "}},\"block\":\"parameters\"}]" |}] +;; diff --git a/src/stan_math_backend/Cpp_Json.mli b/src/stan_math_backend/Cpp_Json.mli index fb17f0b6ee..dfd9b51234 100644 --- a/src/stan_math_backend/Cpp_Json.mli +++ b/src/stan_math_backend/Cpp_Json.mli @@ -1,4 +1,5 @@ open Middle -val out_var_interpolated_json_str : - (string * Expr.Typed.t SizedType.t * Program.io_block) list -> string +val out_var_interpolated_json_str + : (string * Expr.Typed.t SizedType.t * Program.io_block) list + -> string diff --git a/src/stan_math_backend/Expression_gen.ml b/src/stan_math_backend/Expression_gen.ml index 1784705bc9..acd664cdb7 100644 --- a/src/stan_math_backend/Expression_gen.ml +++ b/src/stan_math_backend/Expression_gen.ml @@ -7,28 +7,75 @@ let starts_with prefix s = String.is_prefix ~prefix s let functions_requiring_namespace = String.Set.of_list - [ "e"; "pi"; "log2"; "log10"; "sqrt2"; "not_a_number"; "positive_infinity" - ; "negative_infinity"; "machine_precision"; "abs"; "acos"; "acosh"; "asin" - ; "asinh"; "atan"; "atanh"; "cbrt"; "ceil"; "cos"; "cosh"; "erf"; "erfc" - ; "exp"; "exp2"; "expm1"; "fabs"; "floor"; "lgamma"; "log"; "log1p"; "log2" - ; "log10"; "round"; "sin"; "sinh"; "sqrt"; "tan"; "tanh"; "tgamma"; "trunc" - ; "fdim"; "fmax"; "fmin"; "hypot"; "fma" ] + [ "e" + ; "pi" + ; "log2" + ; "log10" + ; "sqrt2" + ; "not_a_number" + ; "positive_infinity" + ; "negative_infinity" + ; "machine_precision" + ; "abs" + ; "acos" + ; "acosh" + ; "asin" + ; "asinh" + ; "atan" + ; "atanh" + ; "cbrt" + ; "ceil" + ; "cos" + ; "cosh" + ; "erf" + ; "erfc" + ; "exp" + ; "exp2" + ; "expm1" + ; "fabs" + ; "floor" + ; "lgamma" + ; "log" + ; "log1p" + ; "log2" + ; "log10" + ; "round" + ; "sin" + ; "sinh" + ; "sqrt" + ; "tan" + ; "tanh" + ; "tgamma" + ; "trunc" + ; "fdim" + ; "fmax" + ; "fmin" + ; "hypot" + ; "fma" + ] +;; let stan_namespace_qualify f = if Set.mem functions_requiring_namespace f then "stan::math::" ^ f else f +;; (* return true if the types of the two expression are the same *) let types_match e1 e2 = UnsizedType.equal (Expr.Typed.type_of e1) (Expr.Typed.type_of e2) - && UnsizedType.compare_autodifftype (Expr.Typed.adlevel_of e1) + && UnsizedType.compare_autodifftype + (Expr.Typed.adlevel_of e1) (Expr.Typed.adlevel_of e2) = 0 +;; let is_stan_math f = ends_with "__" f || starts_with "stan::math::" f (* retun true if the tpe of the expression is integer or real *) let is_scalar e = - match Expr.Typed.type_of e with UInt | UReal -> true | _ -> false + match Expr.Typed.type_of e with + | UInt | UReal -> true + | _ -> false +;; let is_matrix e = Expr.Typed.type_of e = UMatrix let is_row_vector e = Expr.Typed.type_of e = URowVector @@ -36,42 +83,50 @@ let pretty_print e = Fmt.to_to_string Expr.Typed.pp e let pp_call ppf (name, pp_arg, args) = pf ppf "@[%s(@,%a)@]" name (list ~sep:comma pp_arg) args +;; let rec stantype_prim_str = function | UnsizedType.UInt -> "int" | UArray t -> stantype_prim_str t | _ -> "double" +;; let rec local_scalar ut ad = - match (ut, ad) with + match ut, ad with | UnsizedType.UArray t, _ -> local_scalar t ad | _, UnsizedType.DataOnly | UInt, AutoDiffable -> stantype_prim_str ut | _, AutoDiffable -> "local_scalar_t__" +;; let minus_one e = { e with - Expr.Fixed.pattern= - FunApp (StanLib, Operator.to_string Minus, [e; Expr.Helpers.loop_bottom]) + Expr.Fixed.pattern = + FunApp (StanLib, Operator.to_string Minus, [ e; Expr.Helpers.loop_bottom ]) } +;; -let is_single_index = function Index.Single _ -> true | _ -> false +let is_single_index = function + | Index.Single _ -> true + | _ -> false +;; let dont_need_range_check = function - | Index.Single Expr.Fixed.({pattern= Var id; _}) -> - not (Utils.is_user_ident id) + | Index.Single Expr.Fixed.{ pattern = Var id; _ } -> not (Utils.is_user_ident id) | _ -> false +;; let promote_adtype = List.fold ~f:(fun accum expr -> match Expr.Typed.adlevel_of expr with | AutoDiffable -> AutoDiffable - | _ -> accum ) + | _ -> accum) ~init:UnsizedType.DataOnly +;; let promote_unsizedtype es = let rec fold_type accum mtype = - match (accum, mtype) with + match accum, mtype with | UnsizedType.UReal, _ -> UnsizedType.UReal | _, UnsizedType.UReal -> UReal | UArray t1, UArray t2 -> UArray (fold_type t1 t2) @@ -80,145 +135,136 @@ let promote_unsizedtype es = List.map es ~f:Expr.Typed.type_of |> List.reduce ~f:fold_type |> Option.value ~default:UReal +;; let%expect_test "promote_unsized" = let e mtype = - Expr.{Fixed.pattern= Var "x"; meta= Typed.Meta.{empty with type_= mtype}} + Expr.{ Fixed.pattern = Var "x"; meta = Typed.Meta.{ empty with type_ = mtype } } in let tests = - [[e UInt; e UReal]; [e UReal; e UInt]; [e (UArray UInt); e (UArray UReal)]] + [ [ e UInt; e UReal ]; [ e UReal; e UInt ]; [ e (UArray UInt); e (UArray UReal) ] ] in - print_s - [%sexp (tests |> List.map ~f:promote_unsizedtype : UnsizedType.t list)] ; + print_s [%sexp (tests |> List.map ~f:promote_unsizedtype : UnsizedType.t list)]; [%expect {| (UReal UReal (UArray UReal)) |}] +;; let rec pp_unsizedtype_custom_scalar ppf (scalar, ut) = match ut with | UnsizedType.UInt | UReal -> string ppf scalar - | UArray t -> - pf ppf "std::vector<%a>" pp_unsizedtype_custom_scalar (scalar, t) + | UArray t -> pf ppf "std::vector<%a>" pp_unsizedtype_custom_scalar (scalar, t) | UMatrix -> pf ppf "Eigen::Matrix<%s, -1, -1>" scalar | URowVector -> pf ppf "Eigen::Matrix<%s, 1, -1>" scalar | UVector -> pf ppf "Eigen::Matrix<%s, -1, 1>" scalar | x -> raise_s [%message (x : UnsizedType.t) "not implemented yet"] +;; let pp_unsizedtype_local ppf (adtype, ut) = let s = local_scalar ut adtype in pp_unsizedtype_custom_scalar ppf (s, ut) +;; -let pp_expr_type ppf e = - pp_unsizedtype_local ppf Expr.Typed.(adlevel_of e, type_of e) - -let user_dist_suffices = ["_lpdf"; "_lpmf"; "_log"] +let pp_expr_type ppf e = pp_unsizedtype_local ppf Expr.Typed.(adlevel_of e, type_of e) +let user_dist_suffices = [ "_lpdf"; "_lpmf"; "_log" ] let ends_with_any suffices s = List.exists ~f:(fun suffix -> String.is_suffix ~suffix s) suffices +;; let is_user_dist s = - ends_with_any user_dist_suffices s - && not (ends_with_any ["_cdf_log"; "_ccdf_log"] s) + ends_with_any user_dist_suffices s && not (ends_with_any [ "_cdf_log"; "_ccdf_log" ] s) +;; let is_user_lp s = ends_with "_lp" s let suffix_args f = - if ends_with "_rng" f then ["base_rng__"] - else if ends_with "_lp" f then ["lp__"; "lp_accum__"] + if ends_with "_rng" f + then [ "base_rng__" ] + else if ends_with "_lp" f + then [ "lp__"; "lp_accum__" ] else [] +;; let demangle_propto_name udf f = - if f = "multiply_log" || f = "binomial_coefficient_log" then f - else if Utils.is_propto_distribution f then - Utils.stdlib_distribution_name f ^ "" - else if - Utils.is_distribution_name f || (udf && (is_user_dist f || is_user_lp f)) + if f = "multiply_log" || f = "binomial_coefficient_log" + then f + else if Utils.is_propto_distribution f + then Utils.stdlib_distribution_name f ^ "" + else if Utils.is_distribution_name f || (udf && (is_user_dist f || is_user_lp f)) then f ^ "" else f +;; let fn_renames = List.map - ~f:(fun (k, v) -> (Internal_fun.to_string k, v)) - [ (Internal_fun.FnLength, "stan::math::size") - ; (FnNegInf, "stan::math::negative_infinity") - ; (FnResizeToMatch, "resize_to_match") - ; (FnNaN, "std::numeric_limits::quiet_NaN") ] + ~f:(fun (k, v) -> Internal_fun.to_string k, v) + [ Internal_fun.FnLength, "stan::math::size" + ; FnNegInf, "stan::math::negative_infinity" + ; FnResizeToMatch, "resize_to_match" + ; FnNaN, "std::numeric_limits::quiet_NaN" + ] |> String.Map.of_alist_exn +;; let map_rect_calls = Int.Table.create () let functor_suffix = "_functor__" let reduce_sum_functor_suffix = "_rsfunctor__" let functor_suffix_select hof = - if Stan_math_signatures.is_reduce_sum_fn hof then reduce_sum_functor_suffix + if Stan_math_signatures.is_reduce_sum_fn hof + then reduce_sum_functor_suffix else functor_suffix +;; let rec pp_index ppf = function | Index.All -> pf ppf "index_omni()" | Single e -> pf ppf "index_uni(%a)" pp_expr e | Upfrom e -> pf ppf "index_min(%a)" pp_expr e - | Between (e_low, e_high) -> - pf ppf "index_min_max(%a, %a)" pp_expr e_low pp_expr e_high + | Between (e_low, e_high) -> pf ppf "index_min_max(%a, %a)" pp_expr e_low pp_expr e_high | MultiIndex e -> pf ppf "index_multi(%a)" pp_expr e and pp_indexes ppf = function | [] -> pf ppf "nil_index_list()" - | idx :: idxs -> - pf ppf "@[cons_list(@,%a,@ %a)@]" pp_index idx pp_indexes idxs + | idx :: idxs -> pf ppf "@[cons_list(@,%a,@ %a)@]" pp_index idx pp_indexes idxs and pp_logical_op ppf op lhs rhs = - pf ppf "(primitive_value(@,%a)@ %s@ primitive_value(@,%a))" pp_expr lhs op - pp_expr rhs + pf ppf "(primitive_value(@,%a)@ %s@ primitive_value(@,%a))" pp_expr lhs op pp_expr rhs and pp_unary ppf fm es = pf ppf fm pp_expr (List.hd_exn es) and pp_binary ppf fm es = pf ppf fm pp_expr (first es) pp_expr (second es) - -and pp_binary_f ppf f es = - pf ppf "%s(@,%a,@ %a)" f pp_expr (first es) pp_expr (second es) - +and pp_binary_f ppf f es = pf ppf "%s(@,%a,@ %a)" f pp_expr (first es) pp_expr (second es) and first es = List.nth_exn es 0 and second es = List.nth_exn es 1 and pp_scalar_binary ppf scalar_fmt generic_fmt es = - pp_binary ppf - ( if is_scalar (first es) && is_scalar (second es) then scalar_fmt - else generic_fmt ) + pp_binary + ppf + (if is_scalar (first es) && is_scalar (second es) then scalar_fmt else generic_fmt) es and gen_operator_app = function - | Operator.Plus -> - fun ppf es -> pp_scalar_binary ppf "(%a@ +@ %a)" "add(@,%a,@ %a)" es + | Operator.Plus -> fun ppf es -> pp_scalar_binary ppf "(%a@ +@ %a)" "add(@,%a,@ %a)" es | PMinus -> - fun ppf es -> - pp_unary ppf - (if is_scalar (List.hd_exn es) then "-%a" else "minus(@,%a)") - es + fun ppf es -> + pp_unary ppf (if is_scalar (List.hd_exn es) then "-%a" else "minus(@,%a)") es | PPlus -> fun ppf es -> pp_unary ppf "%a" es | Transpose -> - fun ppf es -> - pp_unary ppf - (if is_scalar (List.hd_exn es) then "%a" else "transpose(@,%a)") - es + fun ppf es -> + pp_unary ppf (if is_scalar (List.hd_exn es) then "%a" else "transpose(@,%a)") es | PNot -> fun ppf es -> pp_unary ppf "logical_negation(@,%a)" es - | Minus -> - fun ppf es -> pp_scalar_binary ppf "(%a@ -@ %a)" "subtract(@,%a,@ %a)" es - | Times -> - fun ppf es -> pp_scalar_binary ppf "(%a@ *@ %a)" "multiply(@,%a,@ %a)" es + | Minus -> fun ppf es -> pp_scalar_binary ppf "(%a@ -@ %a)" "subtract(@,%a,@ %a)" es + | Times -> fun ppf es -> pp_scalar_binary ppf "(%a@ *@ %a)" "multiply(@,%a,@ %a)" es | Divide | IntDivide -> - fun ppf es -> - if - is_matrix (second es) - && (is_matrix (first es) || is_row_vector (first es)) - then pp_binary_f ppf "mdivide_right" es - else pp_scalar_binary ppf "(%a@ /@ %a)" "divide(@,%a,@ %a)" es + fun ppf es -> + if is_matrix (second es) && (is_matrix (first es) || is_row_vector (first es)) + then pp_binary_f ppf "mdivide_right" es + else pp_scalar_binary ppf "(%a@ /@ %a)" "divide(@,%a,@ %a)" es | Modulo -> fun ppf es -> pp_binary_f ppf "modulus" es | LDivide -> fun ppf es -> pp_binary_f ppf "mdivide_left" es - | And | Or -> - raise_s [%message "And/Or should have been converted to an expression"] + | And | Or -> raise_s [%message "And/Or should have been converted to an expression"] | EltTimes -> - fun ppf es -> - pp_scalar_binary ppf "(%a@ *@ %a)" "elt_multiply(@,%a,@ %a)" es + fun ppf es -> pp_scalar_binary ppf "(%a@ *@ %a)" "elt_multiply(@,%a,@ %a)" es | EltDivide -> - fun ppf es -> - pp_scalar_binary ppf "(%a@ /@ %a)" "elt_divide(@,%a,@ %a)" es + fun ppf es -> pp_scalar_binary ppf "(%a@ /@ %a)" "elt_divide(@,%a,@ %a)" es | Pow -> fun ppf es -> pp_binary_f ppf "pow" es | Equals -> fun ppf es -> pp_binary_f ppf "logical_eq" es | NEquals -> fun ppf es -> pp_binary_f ppf "logical_neq" es @@ -229,32 +275,33 @@ and gen_operator_app = function and gen_misc_special_math_app f = match f with - | "lmultiply" -> - Some (fun ppf es -> pp_binary ppf "multiply_log(@,%a,@ %a)" es) + | "lmultiply" -> Some (fun ppf es -> pp_binary ppf "multiply_log(@,%a,@ %a)" es) | "lchoose" -> - Some - (fun ppf es -> pp_binary ppf "binomial_coefficient_log(@,%a,@ %a)" es) + Some (fun ppf es -> pp_binary ppf "binomial_coefficient_log(@,%a,@ %a)" es) | "target" -> Some (fun ppf _ -> pf ppf "get_lp(lp__, lp_accum__)") | "get_lp" -> Some (fun ppf _ -> pf ppf "get_lp(lp__, lp_accum__)") | "max" | "min" -> - Some - (fun ppf es -> - let f = match es with [_; _] -> "std::" ^ f | _ -> f in - pp_call ppf (f, pp_expr, es) ) + Some + (fun ppf es -> + let f = + match es with + | [ _; _ ] -> "std::" ^ f + | _ -> f + in + pp_call ppf (f, pp_expr, es)) | "ceil" -> - let std_prefix_data_scalar f = function - | [ Expr.({ Fixed.meta= - Typed.Meta.({adlevel= DataOnly; type_= UInt | UReal; _}); _ - }) ] -> - "std::" ^ f - | _ -> f - in - Some - (fun ppf es -> - let f = std_prefix_data_scalar f es in - pp_call ppf (f, pp_expr, es) ) + let std_prefix_data_scalar f = function + | [ Expr. + { Fixed.meta = Typed.Meta.{ adlevel = DataOnly; type_ = UInt | UReal; _ }; _ } + ] -> "std::" ^ f + | _ -> f + in + Some + (fun ppf es -> + let f = std_prefix_data_scalar f es in + pp_call ppf (f, pp_expr, es)) | f when Map.mem fn_renames f -> - Some (fun ppf es -> pp_call ppf (Map.find_exn fn_renames f, pp_expr, es)) + Some (fun ppf es -> pp_call ppf (Map.find_exn fn_renames f, pp_expr, es)) | _ -> None and read_data ut ppf es = @@ -263,22 +310,18 @@ and read_data ut ppf es = | UnsizedType.UArray UInt -> "i" | UArray UReal -> "r" | UInt | UReal | UVector | URowVector | UMatrix | UArray _ - |UFun (_, _) - |UMathLibraryFunction -> - raise_s [%message "Can't ReadData of " (ut : UnsizedType.t)] + | UFun (_, _) + | UMathLibraryFunction -> raise_s [%message "Can't ReadData of " (ut : UnsizedType.t)] in pf ppf "context__.vals_%s(%a)" i_or_r pp_expr (List.hd_exn es) (* assumes everything well formed from parser checks *) and gen_fun_app ppf fname es = let default ppf es = - let to_var s = Expr.{Fixed.pattern= Var s; meta= Typed.Meta.empty} in + let to_var s = Expr.{ Fixed.pattern = Var s; meta = Typed.Meta.empty } in let convert_hof_vars = function - | {Expr.Fixed.pattern= Var name; meta= {Expr.Typed.Meta.type_= UFun _; _}} - as e -> - { e with - pattern= FunApp (StanLib, name ^ functor_suffix_select fname, []) - } + | { Expr.Fixed.pattern = Var name; meta = { Expr.Typed.Meta.type_ = UFun _; _ } } as + e -> { e with pattern = FunApp (StanLib, name ^ functor_suffix_select fname, []) } | e -> e in let converted_es = List.map ~f:convert_hof_vars es in @@ -294,60 +337,67 @@ and gen_fun_app ppf fname es = overloads. *) let fname, args = - match (is_hof_call, fname, converted_es @ extra) with + match is_hof_call, fname, converted_es @ extra with | true, "algebra_solver", f :: x :: y :: dat :: datint :: tl - |true, "algebra_solver_newton", f :: x :: y :: dat :: datint :: tl -> - (fname, f :: x :: y :: dat :: datint :: msgs :: tl) + | true, "algebra_solver_newton", f :: x :: y :: dat :: datint :: tl -> + fname, f :: x :: y :: dat :: datint :: msgs :: tl | true, "integrate_1d", f :: a :: b :: theta :: x_r :: x_i :: tl -> - (fname, f :: a :: b :: theta :: x_r :: x_i :: msgs :: tl) - | ( true - , "integrate_ode_bdf" - , f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl ) - |( true - , "integrate_ode_adams" - , f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl ) - |( true - , "integrate_ode_rk45" - , f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl ) -> - (fname, f :: y0 :: t0 :: ts :: theta :: x :: x_int :: msgs :: tl) - | true, x, {pattern= FunApp (_, f, _); _} :: grainsize :: container :: tl + fname, f :: a :: b :: theta :: x_r :: x_i :: msgs :: tl + | true, "integrate_ode_bdf", f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl + | true, "integrate_ode_adams", f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl + | true, "integrate_ode_rk45", f :: y0 :: t0 :: ts :: theta :: x :: x_int :: tl -> + fname, f :: y0 :: t0 :: ts :: theta :: x :: x_int :: msgs :: tl + | true, x, { pattern = FunApp (_, f, _); _ } :: grainsize :: container :: tl when Stan_math_signatures.is_reduce_sum_fn x -> - (strf "%s<%s>" fname f, grainsize :: container :: msgs :: tl) - | true, "map_rect", {pattern= FunApp (_, f, _); _} :: tl -> - let next_map_rect_id = Hashtbl.length map_rect_calls + 1 in - Hashtbl.add_exn map_rect_calls ~key:next_map_rect_id ~data:f ; - (strf "%s<%d, %s>" fname next_map_rect_id f, tl @ [msgs]) - | true, _, args -> (fname, args @ [msgs]) - | false, _, args -> (fname, args) + strf "%s<%s>" fname f, grainsize :: container :: msgs :: tl + | true, "map_rect", { pattern = FunApp (_, f, _); _ } :: tl -> + let next_map_rect_id = Hashtbl.length map_rect_calls + 1 in + Hashtbl.add_exn map_rect_calls ~key:next_map_rect_id ~data:f; + strf "%s<%d, %s>" fname next_map_rect_id f, tl @ [ msgs ] + | true, _, args -> fname, args @ [ msgs ] + | false, _, args -> fname, args in let fname = stan_namespace_qualify fname |> demangle_propto_name false in pp_call ppf (fname, pp_expr, args) in let pp = [ Option.map ~f:gen_operator_app (Operator.of_string_opt fname) - ; gen_misc_special_math_app fname ] - |> List.filter_opt |> List.hd |> Option.value ~default + ; gen_misc_special_math_app fname + ] + |> List.filter_opt + |> List.hd + |> Option.value ~default in pf ppf "@[%a@]" pp es and pp_constrain_funapp constrain_or_un_str ppf = function - | var :: {Expr.Fixed.pattern= Lit (Str, constraint_flavor); _} :: args -> - pf ppf "@[stan::math::%s_%s(@,%a@])" constraint_flavor - constrain_or_un_str (list ~sep:comma pp_expr) (var :: args) + | var :: { Expr.Fixed.pattern = Lit (Str, constraint_flavor); _ } :: args -> + pf + ppf + "@[stan::math::%s_%s(@,%a@])" + constraint_flavor + constrain_or_un_str + (list ~sep:comma pp_expr) + (var :: args) | es -> raise_s [%message "Bad constraint " (es : Expr.Typed.t list)] and pp_user_defined_fun ppf (f, es) = - let extra_args = suffix_args f @ ["pstream__"] in + let extra_args = suffix_args f @ [ "pstream__" ] in let sep = if List.is_empty es then "" else ", " in - pf ppf "@[%s(@,%a%s)@]" + pf + ppf + "@[%s(@,%a%s)@]" (demangle_propto_name true f) - (list ~sep:comma pp_expr) es + (list ~sep:comma pp_expr) + es (sep ^ String.concat ~sep:", " extra_args) and pp_compiler_internal_fn ut f ppf es = let pp_array_literal ppf es = let pp_add_method ppf () = pf ppf ")@,.add(" in - pf ppf "stan::math::array_builder<%a>()@,.add(%a)@,.array()" + pf + ppf + "stan::math::array_builder<%a>()@,.add(%a)@,.array()" pp_unsizedtype_local (promote_adtype es, promote_unsizedtype es) (list ~sep:pp_add_method pp_expr) @@ -355,24 +405,21 @@ and pp_compiler_internal_fn ut f ppf es = in match Internal_fun.of_string_opt f with | Some FnMakeArray -> pp_array_literal ppf es - | Some FnMakeRowVec -> ( - match ut with + | Some FnMakeRowVec -> + (match ut with | UnsizedType.URowVector -> - pf ppf "stan::math::to_row_vector(@,%a)" pp_array_literal es + pf ppf "stan::math::to_row_vector(@,%a)" pp_array_literal es | UMatrix -> pf ppf "stan::math::to_matrix(@,%a)" pp_array_literal es | _ -> - raise_s - [%message - "Unexpected type for row vector literal" (ut : UnsizedType.t)] ) + raise_s [%message "Unexpected type for row vector literal" (ut : UnsizedType.t)]) | Some FnConstrain -> pp_constrain_funapp "constrain" ppf es | Some FnUnconstrain -> pp_constrain_funapp "free" ppf es | Some FnReadData -> read_data ut ppf es - | Some FnReadParam -> ( - match es with - | {Expr.Fixed.pattern= Lit (Str, base_type); _} :: dims -> - pf ppf "@[in__.%s(@,%a)@]" base_type (list ~sep:comma pp_expr) - dims - | _ -> raise_s [%message "emit ReadParam with " (es : Expr.Typed.t list)] ) + | Some FnReadParam -> + (match es with + | { Expr.Fixed.pattern = Lit (Str, base_type); _ } :: dims -> + pf ppf "@[in__.%s(@,%a)@]" base_type (list ~sep:comma pp_expr) dims + | _ -> raise_s [%message "emit ReadParam with " (es : Expr.Typed.t list)]) | _ -> gen_fun_app ppf f es and pp_indexed ppf (vident, indices, pretty) = @@ -382,121 +429,139 @@ and pp_indexed_simple ppf (obj, idcs) = let idx_minus_one = function | Index.Single e -> minus_one e | MultiIndex e | Between (e, _) | Upfrom e -> - raise_s - [%message - "No non-Single indices allowed" ~obj - (idcs : Expr.Typed.t Index.t list) - (Expr.Typed.loc_of e : Location_span.t)] + raise_s + [%message + "No non-Single indices allowed" + ~obj + (idcs : Expr.Typed.t Index.t list) + (Expr.Typed.loc_of e : Location_span.t)] | All -> - raise_s - [%message - "No non-Single indices allowed" ~obj - (idcs : Expr.Typed.t Index.t list)] + raise_s + [%message "No non-Single indices allowed" ~obj (idcs : Expr.Typed.t Index.t list)] in - pf ppf "%s%a" obj + pf + ppf + "%s%a" + obj (fun ppf idcs -> match idcs with | [] -> () - | idcs -> pf ppf "[%a]" (list ~sep:(const string "][") pp_expr) idcs ) + | idcs -> pf ppf "[%a]" (list ~sep:(const string "][") pp_expr) idcs) (List.map ~f:idx_minus_one idcs) -and pp_expr ppf Expr.Fixed.({pattern; meta} as e) = +and pp_expr ppf Expr.Fixed.({ pattern; meta } as e) = match pattern with | Var s -> pf ppf "%s" s | Lit (Str, s) -> pf ppf "%S" s | Lit (_, s) -> pf ppf "%s" s | FunApp (StanLib, f, es) -> gen_fun_app ppf f es | FunApp (CompilerInternal, f, es) -> - pp_compiler_internal_fn meta.type_ (stan_namespace_qualify f) ppf es + pp_compiler_internal_fn meta.type_ (stan_namespace_qualify f) ppf es | FunApp (UserDefined, f, es) -> pp_user_defined_fun ppf (f, es) | EAnd (e1, e2) -> pp_logical_op ppf "&&" e1 e2 | EOr (e1, e2) -> pp_logical_op ppf "||" e1 e2 | TernaryIf (ec, et, ef) -> - let promoted ppf (t, e) = - pf ppf "stan::math::promote_scalar<%s>(%a)" - Expr.Typed.(local_scalar (type_of t) (adlevel_of t)) - pp_expr e - in - let tform ppf = pf ppf "(@[@,%a@ ?@ %a@ :@ %a@])" in - if types_match et ef then tform ppf pp_expr ec pp_expr et pp_expr ef - else tform ppf pp_expr ec promoted (e, et) promoted (e, ef) + let promoted ppf (t, e) = + pf + ppf + "stan::math::promote_scalar<%s>(%a)" + Expr.Typed.(local_scalar (type_of t) (adlevel_of t)) + pp_expr + e + in + let tform ppf = pf ppf "(@[@,%a@ ?@ %a@ :@ %a@])" in + if types_match et ef + then tform ppf pp_expr ec pp_expr et pp_expr ef + else tform ppf pp_expr ec promoted (e, et) promoted (e, ef) | Indexed (e, []) -> pp_expr ppf e - | Indexed (e, idx) -> ( - match e.pattern with + | Indexed (e, idx) -> + (match e.pattern with | FunApp (CompilerInternal, f, _) - when Some Internal_fun.FnReadParam = Internal_fun.of_string_opt f -> - pp_expr ppf e + when Some Internal_fun.FnReadParam = Internal_fun.of_string_opt f -> pp_expr ppf e | FunApp (CompilerInternal, f, _) when Some Internal_fun.FnReadData = Internal_fun.of_string_opt f -> - pp_indexed_simple ppf (strf "%a" pp_expr e, idx) + pp_indexed_simple ppf (strf "%a" pp_expr e, idx) | _ when List.for_all ~f:dont_need_range_check idx - && not (UnsizedType.is_indexing_matrix (Expr.Typed.type_of e, idx)) - -> - pp_indexed_simple ppf (strf "%a" pp_expr e, idx) - | _ -> pp_indexed ppf (strf "%a" pp_expr e, idx, pretty_print e) ) + && not (UnsizedType.is_indexing_matrix (Expr.Typed.type_of e, idx)) -> + pp_indexed_simple ppf (strf "%a" pp_expr e, idx) + | _ -> pp_indexed ppf (strf "%a" pp_expr e, idx, pretty_print e)) +;; (* these functions are just for testing *) let dummy_locate pattern = Expr.( Fixed. { pattern - ; meta= - Typed.Meta.{type_= UInt; adlevel= DataOnly; loc= Location_span.empty} + ; meta = Typed.Meta.{ type_ = UInt; adlevel = DataOnly; loc = Location_span.empty } }) +;; let pp_unlocated e = strf "%a" pp_expr (dummy_locate e) let%expect_test "pp_expr1" = - printf "%s" (pp_unlocated (Var "a")) ; + printf "%s" (pp_unlocated (Var "a")); [%expect {| a |}] +;; let%expect_test "pp_expr2" = - printf "%s" (pp_unlocated (Lit (Str, "b"))) ; + printf "%s" (pp_unlocated (Lit (Str, "b"))); [%expect {| "b" |}] +;; let%expect_test "pp_expr3" = - printf "%s" (pp_unlocated (Lit (Int, "112"))) ; + printf "%s" (pp_unlocated (Lit (Int, "112"))); [%expect {| 112 |}] +;; let%expect_test "pp_expr4" = - printf "%s" (pp_unlocated (Lit (Int, "112"))) ; + printf "%s" (pp_unlocated (Lit (Int, "112"))); [%expect {| 112 |}] +;; let%expect_test "pp_expr5" = - printf "%s" (pp_unlocated (FunApp (StanLib, "pi", []))) ; + printf "%s" (pp_unlocated (FunApp (StanLib, "pi", []))); [%expect {| stan::math::pi() |}] +;; let%expect_test "pp_expr6" = - printf "%s" - (pp_unlocated (FunApp (StanLib, "sqrt", [dummy_locate (Lit (Int, "123"))]))) ; + printf + "%s" + (pp_unlocated (FunApp (StanLib, "sqrt", [ dummy_locate (Lit (Int, "123")) ]))); [%expect {| stan::math::sqrt(123) |}] +;; let%expect_test "pp_expr7" = - printf "%s" + printf + "%s" (pp_unlocated (FunApp ( StanLib , "atan" - , [dummy_locate (Lit (Int, "123")); dummy_locate (Lit (Real, "1.2"))] - ))) ; + , [ dummy_locate (Lit (Int, "123")); dummy_locate (Lit (Real, "1.2")) ] ))); [%expect {| stan::math::atan(123, 1.2) |}] +;; let%expect_test "pp_expr9" = - printf "%s" + printf + "%s" (pp_unlocated (TernaryIf ( dummy_locate (Lit (Int, "1")) , dummy_locate (Lit (Real, "1.2")) - , dummy_locate (Lit (Real, "2.3")) ))) ; + , dummy_locate (Lit (Real, "2.3")) ))); [%expect {| (1 ? 1.2 : 2.3) |}] +;; let%expect_test "pp_expr10" = - printf "%s" (pp_unlocated (Indexed (dummy_locate (Var "a"), [All]))) ; + printf "%s" (pp_unlocated (Indexed (dummy_locate (Var "a"), [ All ]))); [%expect {| rvalue(a, cons_list(index_omni(), nil_index_list()), "a") |}] +;; let%expect_test "pp_expr11" = - printf "%s" + printf + "%s" (pp_unlocated - (FunApp (UserDefined, "poisson_rng", [dummy_locate (Lit (Int, "123"))]))) ; + (FunApp (UserDefined, "poisson_rng", [ dummy_locate (Lit (Int, "123")) ]))); [%expect {| poisson_rng(123, base_rng__, pstream__) |}] +;; diff --git a/src/stan_math_backend/Locations.ml b/src/stan_math_backend/Locations.ml index fd6621d410..25872d7d7b 100644 --- a/src/stan_math_backend/Locations.ml +++ b/src/stan_math_backend/Locations.ml @@ -8,42 +8,47 @@ let no_span_num = 0 let prepare_prog (mir : Program.Typed.t) : Program.Numbered.t * state_t = let label_to_location = Int.Table.create () in let location_to_label = Hashtbl.create (module Location_span) in - Hashtbl.set label_to_location ~key:no_span_num ~data:Location_span.empty ; - Hashtbl.set location_to_label ~key:Location_span.empty ~data:no_span_num ; - let rec number_locations_stmt ({pattern; meta} : Stmt.Located.t) : - Stmt.Numbered.t = + Hashtbl.set label_to_location ~key:no_span_num ~data:Location_span.empty; + Hashtbl.set location_to_label ~key:Location_span.empty ~data:no_span_num; + let rec number_locations_stmt ({ pattern; meta } : Stmt.Located.t) : Stmt.Numbered.t = let pattern = Stmt.Fixed.Pattern.map Fn.id number_locations_stmt pattern in match Hashtbl.find location_to_label meta with | Some i -> - let meta = Stmt.Numbered.Meta.from_int i in - {meta; pattern} + let meta = Stmt.Numbered.Meta.from_int i in + { meta; pattern } | None -> - let new_label = Hashtbl.length label_to_location in - Hashtbl.set label_to_location ~key:new_label ~data:meta ; - Hashtbl.set location_to_label ~key:meta ~data:new_label ; - {pattern; meta= new_label} + let new_label = Hashtbl.length label_to_location in + Hashtbl.set label_to_location ~key:new_label ~data:meta; + Hashtbl.set location_to_label ~key:meta ~data:new_label; + { pattern; meta = new_label } in let mir = Program.map Fn.id number_locations_stmt mir in let location_list = - List.map ~f:snd + List.map + ~f:snd (List.sort ~compare:(fun x y -> compare_int (fst x) (fst y)) (Hashtbl.to_alist label_to_location)) in - (mir, location_list) + mir, location_list +;; let pp_globals ppf location_list = let location_list = " (found before start of program)" - :: ( List.filter ~f:(fun x -> x <> Location_span.empty) location_list - |> List.map ~f:(fun x -> " (in " ^ Location_span.to_string x ^ ")") ) + :: (List.filter ~f:(fun x -> x <> Location_span.empty) location_list + |> List.map ~f:(fun x -> " (in " ^ Location_span.to_string x ^ ")")) in - Fmt.pf ppf + Fmt.pf + ppf "@ static int current_statement__ = 0;@ static const std::vector \ locations_array__ = {@[%a@]};@ " Fmt.(list ~sep:comma (fmt "%S")) location_list +;; let pp_smeta ppf location_num = - if location_num = no_span_num then () + if location_num = no_span_num + then () else Fmt.pf ppf "current_statement__ = %d;@;" location_num +;; diff --git a/src/stan_math_backend/Stan_math_code_gen.ml b/src/stan_math_backend/Stan_math_code_gen.ml index d81c8319d3..9c24b6e17f 100644 --- a/src/stan_math_backend/Stan_math_code_gen.ml +++ b/src/stan_math_backend/Stan_math_code_gen.ml @@ -27,9 +27,11 @@ let stanc_args_to_print = not String.(is_suffix ~suffix:".stan" x || is_prefix ~prefix:"--o" x) in (* Ignore the "--o" arg, the stan file and the binary name (bin/stanc). *) - Array.to_list Sys.argv |> List.tl_exn + Array.to_list Sys.argv + |> List.tl_exn |> List.filter ~f:sans_model_and_hpp_paths |> String.concat ~sep:" " +;; let pp_unused = fmt "(void) %s; // suppress unused var warning@ " @@ -38,22 +40,25 @@ let pp_unused = fmt "(void) %s; // suppress unused var warning@ " @param fname Name of the function. *) let pp_function__ ppf (prog_name, fname) = - pf ppf "static const char* function__ = %S;@ " - (strf "%s_namespace::%s" prog_name fname) ; + pf ppf "static const char* function__ = %S;@ " (strf "%s_namespace::%s" prog_name fname); pp_unused ppf "function__" +;; (** Print the body of exception handling for functions *) let pp_located ppf _ = - pf ppf + pf + ppf {|stan::lang::rethrow_located(e, locations_array__[current_statement__]); // Next line prevents compiler griping about no return throw std::runtime_error("*** IF YOU SEE THIS, PLEASE REPORT A BUG ***"); |} +;; (** Detect if argument requires C++ template *) let arg_needs_template = function | UnsizedType.DataOnly, _, _ -> false | _, _, t when UnsizedType.contains_int t -> false | _ -> true +;; (** Print template arguments for C++ functions that need templates @param args A pack of `Program.fun_arg_decl` containing functions to detect templates. @@ -63,13 +68,17 @@ let maybe_templated_arg_types (args : Program.fun_arg_decl) = List.mapi args ~f:(fun i a -> match arg_needs_template a with | true -> Some (sprintf "T%d__" i) - | false -> None ) + | false -> None) +;; let%expect_test "arg types templated correctly" = - [(AutoDiffable, "xreal", UReal); (DataOnly, "yint", UInt)] - |> maybe_templated_arg_types |> List.filter_opt |> String.concat ~sep:"," - |> print_endline ; + [ AutoDiffable, "xreal", UReal; DataOnly, "yint", UInt ] + |> maybe_templated_arg_types + |> List.filter_opt + |> String.concat ~sep:"," + |> print_endline; [%expect {| T0__ |}] +;; (** Print the code for promoting stan real types @param ppf A pretty printer @@ -79,19 +88,20 @@ let pp_promoted_scalar ppf args = match args with | [] -> pf ppf "double" | _ -> - let rec promote_args_chunked ppf args = - let go ppf tl = - match tl with [] -> () | _ -> pf ppf ", %a" promote_args_chunked tl - in - match args with - | [] -> pf ppf "double" - | hd :: tl -> - pf ppf "stan::promote_args_t<%a%a>" (list ~sep:comma string) hd go - tl + let rec promote_args_chunked ppf args = + let go ppf tl = + match tl with + | [] -> () + | _ -> pf ppf ", %a" promote_args_chunked tl in - promote_args_chunked ppf - List.( - chunks_of ~length:5 (filter_opt (maybe_templated_arg_types args))) + match args with + | [] -> pf ppf "double" + | hd :: tl -> pf ppf "stan::promote_args_t<%a%a>" (list ~sep:comma string) hd go tl + in + promote_args_chunked + ppf + List.(chunks_of ~length:5 (filter_opt (maybe_templated_arg_types args))) +;; (** Pretty-prints a function's return-type, taking into account templated argument promotion.*) @@ -99,9 +109,10 @@ let pp_returntype ppf arg_types rt = let scalar = strf "%a" pp_promoted_scalar arg_types in match rt with | Some ut when UnsizedType.contains_int ut -> - pf ppf "%a@," pp_unsizedtype_custom_scalar ("int", 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@," +;; (** [pp_located_error ppf (pp_body_block, body_block, err_msg)] surrounds [body_block] with a C++ try-catch that will rethrow the error with the proper source location @@ -111,9 +122,10 @@ let pp_returntype ppf arg_types rt = @param body A C++ scoped body block surrounded by squiggly braces. *) let pp_located_error ppf (pp_body_block, body) = - pf ppf "@ try %a" pp_body_block body ; - string ppf " catch (const std::exception& e) " ; + pf ppf "@ try %a" pp_body_block body; + string ppf " catch (const std::exception& e) "; pp_block ppf (pp_located, ()) +;; (** Print the type of an object. @param ppf A pretty printer @@ -128,12 +140,14 @@ let pp_arg ppf (custom_scalar_opt, (_, name, ut)) = | None -> stantype_prim_str ut in pf ppf "const %a& %s" pp_unsizedtype_custom_scalar (scalar, ut) name +;; (** [pp_located_error_b] automatically adds a Block wrapper *) let pp_located_error_b ppf body_stmts = - pp_located_error ppf - ( pp_statement - , Stmt.Fixed.{pattern= Block body_stmts; meta= Locations.no_span_num} ) + pp_located_error + ppf + (pp_statement, Stmt.Fixed.{ pattern = Block body_stmts; meta = Locations.no_span_num }) +;; let typename = ( ^ ) "typename " @@ -143,106 +157,119 @@ let typename = ( ^ ) "typename " let get_templates_and_args fdargs = let argtypetemplates = maybe_templated_arg_types fdargs in ( List.filter_opt argtypetemplates - , List.map - ~f:(fun a -> strf "%a" pp_arg a) - (List.zip_exn argtypetemplates fdargs) ) + , List.map ~f:(fun a -> strf "%a" pp_arg a) (List.zip_exn argtypetemplates fdargs) ) +;; (** Print the C++ template parameter decleration before a function. @param ppf A pretty printer. *) let pp_template_decorator ppf = function | [] -> () - | templates -> - pf ppf "@[template <%a>@]@ " (list ~sep:comma string) templates + | templates -> pf ppf "@[template <%a>@]@ " (list ~sep:comma string) templates +;; (** Print the C++ function definition. @param ppf A pretty printer Refactor this please - one idea might be to have different functions for printing user defined distributions vs rngs vs regular functions. *) -let pp_fun_def ppf Program.({fdrt; fdname; fdargs; fdbody; _}) - funs_used_in_reduce_sum = +let pp_fun_def ppf Program.{ fdrt; fdname; fdargs; fdbody; _ } funs_used_in_reduce_sum = let is_lp = is_user_lp fdname in let is_dist = is_user_dist fdname in let is_rng = String.is_suffix fdname ~suffix:"_rng" in let extra, extra_templates = - if is_lp then (["lp__"; "lp_accum__"], ["T_lp__"; "T_lp_accum__"]) - else if is_rng then (["base_rng__"], ["RNG"]) - else ([], []) + if is_lp + then [ "lp__"; "lp_accum__" ], [ "T_lp__"; "T_lp_accum__" ] + else if is_rng + then [ "base_rng__" ], [ "RNG" ] + else [], [] in let mk_extra_args templates args = List.map ~f:(fun (t, v) -> t ^ "& " ^ v) (List.zip_exn templates args) in let argtypetemplates, args = get_templates_and_args fdargs in - let pp_body ppf (Stmt.Fixed.({pattern; _}) as fdbody) = + let pp_body ppf (Stmt.Fixed.{ pattern; _ } as fdbody) = let text = pf ppf "%s@;" in - pf ppf "@[using local_scalar_t__ = %a;@]@," pp_promoted_scalar fdargs ; - if not (is_dist || is_lp) then ( - text "const static bool propto__ = true;" ; - text "(void) propto__;" ) ; - text - "local_scalar_t__ DUMMY_VAR__(std::numeric_limits::quiet_NaN());" ; - pp_unused ppf "DUMMY_VAR__" ; + pf ppf "@[using local_scalar_t__ = %a;@]@," pp_promoted_scalar fdargs; + if not (is_dist || is_lp) + then ( + text "const static bool propto__ = true;"; + text "(void) propto__;"); + text "local_scalar_t__ DUMMY_VAR__(std::numeric_limits::quiet_NaN());"; + pp_unused ppf "DUMMY_VAR__"; let blocked_fdbody = match pattern with - | SList stmts -> {fdbody with pattern= Block stmts} + | SList stmts -> { fdbody with pattern = Block stmts } | Block _ -> fdbody - | _ -> {fdbody with pattern= Block [fdbody]} + | _ -> { fdbody with pattern = Block [ fdbody ] } in - pp_located_error ppf (pp_statement, blocked_fdbody) ; + pp_located_error ppf (pp_statement, blocked_fdbody); pf ppf "@ " in let templates = - (if is_dist || is_lp then ["bool propto__"] else []) + (if is_dist || is_lp then [ "bool propto__" ] else []) @ List.(map ~f:typename (argtypetemplates @ extra_templates)) in let pp_sig ppf name = - pp_template_decorator ppf templates ; - pp_returntype ppf fdargs fdrt ; + pp_template_decorator ppf templates; + pp_returntype ppf fdargs fdrt; let arg_strs = - args @ mk_extra_args extra_templates extra @ ["std::ostream* pstream__"] + args @ mk_extra_args extra_templates extra @ [ "std::ostream* pstream__" ] in pf ppf "%s(@[%a@]) " name (list ~sep:comma string) arg_strs in let pp_sig_rs ppf name = - if is_dist then pp_template_decorator ppf (List.tl_exn templates) - else pp_template_decorator ppf templates ; - pp_returntype ppf fdargs fdrt ; + if is_dist + then pp_template_decorator ppf (List.tl_exn templates) + else pp_template_decorator ppf templates; + pp_returntype ppf fdargs fdrt; let first_three, rest = List.split_n args 3 in let arg_strs = first_three - @ ["std::ostream* pstream__"] + @ [ "std::ostream* pstream__" ] @ rest @ mk_extra_args extra_templates extra in pf ppf "%s(@[%a@]) " name (list ~sep:comma string) arg_strs in - pp_sig ppf fdname ; + pp_sig ppf fdname; match Stmt.Fixed.(fdbody.pattern) with | Skip -> pf ppf ";@ " - | _ -> ( - pp_block ppf (pp_body, fdbody) ; - pf ppf "@,@,struct %s%s {@,%a const @,{@,return %a;@,}@,};@," fdname - functor_suffix pp_sig "operator()" pp_call_str - ( (if is_dist || is_lp then fdname ^ "" else fdname) - , List.map ~f:(fun (_, name, _) -> name) fdargs @ extra @ ["pstream__"] - ) ; - if String.Set.mem funs_used_in_reduce_sum fdname then - (* Produces the reduce_sum functors that has the pstream argument - as the third and not last argument *) - match fdargs with - | (_, slice, _) :: (_, start, _) :: (_, end_, _) :: rest -> - pf ppf "@,@,struct %s%s {@,%a const @,{@,return %a;@,}@,};@," - fdname reduce_sum_functor_suffix pp_sig_rs "operator()" - pp_call_str - ( (if is_dist then fdname ^ "" else fdname) - , slice :: (start ^ " + 1") :: (end_ ^ " + 1") - :: List.map ~f:(fun (_, name, _) -> name) rest - @ extra @ ["pstream__"] ) - | _ -> - raise_s - [%message - "Ill-formed reduce_sum call! This is bug in the compiler."] ) + | _ -> + pp_block ppf (pp_body, fdbody); + pf + ppf + "@,@,struct %s%s {@,%a const @,{@,return %a;@,}@,};@," + fdname + functor_suffix + pp_sig + "operator()" + pp_call_str + ( (if is_dist || is_lp then fdname ^ "" else fdname) + , List.map ~f:(fun (_, name, _) -> name) fdargs @ extra @ [ "pstream__" ] ); + if String.Set.mem funs_used_in_reduce_sum fdname + then ( + (* Produces the reduce_sum functors that has the pstream argument + as the third and not last argument *) + match fdargs with + | (_, slice, _) :: (_, start, _) :: (_, end_, _) :: rest -> + pf + ppf + "@,@,struct %s%s {@,%a const @,{@,return %a;@,}@,};@," + fdname + reduce_sum_functor_suffix + pp_sig_rs + "operator()" + pp_call_str + ( (if is_dist then fdname ^ "" else fdname) + , (slice + :: (start ^ " + 1") + :: (end_ ^ " + 1") + :: List.map ~f:(fun (_, name, _) -> name) rest) + @ extra + @ [ "pstream__" ] ) + | _ -> raise_s [%message "Ill-formed reduce_sum call! This is bug in the compiler."]) +;; let version = "// Code generated by %%NAME%% %%VERSION%%" let includes = "#include " @@ -253,13 +280,18 @@ let includes = "#include " @param st The SizedType of the object. *) let pp_validate_data ppf (name, st) = - if String.is_suffix ~suffix:"__" name then () + if String.is_suffix ~suffix:"__" name + then () else - pf ppf "@[context__.validate_dims(@,%S,@,%S,@,%S,@,%a);@]@ " - "data initialization" name + pf + ppf + "@[context__.validate_dims(@,%S,@,%S,@,%S,@,%a);@]@ " + "data initialization" + name (stantype_prim_str (SizedType.to_unsized st)) pp_call ("context__.to_vec", pp_expr, SizedType.get_dims st) +;; (** Print the constructor of the model class. Read in data steps: @@ -270,86 +302,95 @@ let pp_validate_data ppf (name, st) = *) let pp_ctor ppf p = let params = - [ "stan::io::var_context& context__"; "unsigned int random_seed__ = 0" - ; "std::ostream* pstream__ = nullptr" ] - in - pf ppf "%s(@[%a) : model_base_crtp(0) @]" p.Program.prog_name - (list ~sep:comma string) params ; + [ "stan::io::var_context& context__" + ; "unsigned int random_seed__ = 0" + ; "std::ostream* pstream__ = nullptr" + ] + in + pf + ppf + "%s(@[%a) : model_base_crtp(0) @]" + p.Program.prog_name + (list ~sep:comma string) + params; let pp_mul ppf () = pf ppf " * " in let pp_num_param ppf (checks, dims) = - if List.length checks > 0 then ( - list ~sep:cut pp_statement ppf checks ; - cut ppf () ) ; + if List.length checks > 0 + then ( + list ~sep:cut pp_statement ppf checks; + cut ppf ()); pf ppf "num_params_r__ += %a;" (list ~sep:pp_mul pp_expr) dims in let get_param_st = function | ( decl_id - , { Program.out_block= Parameters - ; out_unconstrained_st= st - ; out_constrained_st= cst - ; out_trans= tr } ) -> ( - let meta = - p.log_prob - |> List.find ~f:(function - | {Stmt.Fixed.pattern= Decl {decl_id= id; _}; _} - when id = decl_id -> - true - | _ -> false ) - |> Option.map ~f:(fun x -> x.meta) - |> Option.value ~default:Stmt.Numbered.Meta.empty - in - let dims_check = Transform_Mir.validate_sized decl_id meta (Some tr) in - match SizedType.get_dims st with - | [] -> Some (dims_check cst, [Expr.Helpers.loop_bottom]) - | ls -> Some (dims_check cst, ls) ) + , { Program.out_block = Parameters + ; out_unconstrained_st = st + ; out_constrained_st = cst + ; out_trans = tr + } ) -> + let meta = + p.log_prob + |> List.find ~f:(function + | { Stmt.Fixed.pattern = Decl { decl_id = id; _ }; _ } when id = decl_id -> + true + | _ -> false) + |> Option.map ~f:(fun x -> x.meta) + |> Option.value ~default:Stmt.Numbered.Meta.empty + in + let dims_check = Transform_Mir.validate_sized decl_id meta (Some tr) in + (match SizedType.get_dims st with + | [] -> Some (dims_check cst, [ Expr.Helpers.loop_bottom ]) + | ls -> Some (dims_check cst, ls)) | _ -> None in let data_idents = List.map ~f:fst p.input_vars |> String.Set.of_list in - let pp_stmt_topdecl_size_only ppf (Stmt.Fixed.({pattern; meta}) as s) = + let pp_stmt_topdecl_size_only ppf (Stmt.Fixed.{ pattern; meta } as s) = match pattern with - | Decl {decl_id; decl_type; _} -> ( - match decl_type with + | Decl { decl_id; decl_type; _ } -> + (match decl_type with | Sized st -> - Locations.pp_smeta ppf meta ; - if Set.mem data_idents decl_id then pp_validate_data ppf (decl_id, st) ; - pp_set_size ppf (decl_id, st, DataOnly) - | Unsized _ -> () ) + Locations.pp_smeta ppf meta; + if Set.mem data_idents decl_id then pp_validate_data ppf (decl_id, st); + pp_set_size ppf (decl_id, st, DataOnly) + | Unsized _ -> ()) | _ -> pp_statement ppf s in - pp_block ppf - ( (fun ppf {Program.prog_name; prepare_data; output_vars; _} -> - pf ppf "using local_scalar_t__ = double ;@ " ; - pf ppf "boost::ecuyer1988 base_rng__ = @ " ; - pf ppf " stan::services::util::create_rng(random_seed__, 0);@ " ; - pp_unused ppf "base_rng__" ; - pp_function__ ppf (prog_name, prog_name) ; - pf ppf - "local_scalar_t__ \ - DUMMY_VAR__(std::numeric_limits::quiet_NaN());@ " ; - pp_unused ppf "DUMMY_VAR__" ; - pp_located_error ppf - (pp_block, (list ~sep:cut pp_stmt_topdecl_size_only, prepare_data)) ; - cut ppf () ; - pf ppf "num_params_r__ = 0U;@ " ; - pp_located_error ppf + pp_block + ppf + ( (fun ppf { Program.prog_name; prepare_data; output_vars; _ } -> + pf ppf "using local_scalar_t__ = double ;@ "; + pf ppf "boost::ecuyer1988 base_rng__ = @ "; + pf ppf " stan::services::util::create_rng(random_seed__, 0);@ "; + pp_unused ppf "base_rng__"; + pp_function__ ppf (prog_name, prog_name); + pf ppf "local_scalar_t__ DUMMY_VAR__(std::numeric_limits::quiet_NaN());@ "; + pp_unused ppf "DUMMY_VAR__"; + pp_located_error + ppf + (pp_block, (list ~sep:cut pp_stmt_topdecl_size_only, prepare_data)); + cut ppf (); + pf ppf "num_params_r__ = 0U;@ "; + pp_located_error + ppf ( pp_block - , ( list ~sep:cut pp_num_param - , List.filter_map ~f:get_param_st output_vars ) ) ) + , (list ~sep:cut pp_num_param, List.filter_map ~f:get_param_st output_vars) )) , p ) +;; -let rec top_level_decls Stmt.Fixed.({pattern; _}) = +let rec top_level_decls Stmt.Fixed.{ pattern; _ } = match pattern with - | Decl d -> - [Some (d.decl_id, Type.to_unsized d.decl_type, UnsizedType.DataOnly)] + | Decl d -> [ Some (d.decl_id, Type.to_unsized d.decl_type, UnsizedType.DataOnly) ] | SList stmts -> List.concat_map ~f:top_level_decls stmts - | _ -> [None] + | _ -> [ None ] +;; (** Print the private data members of the model class *) -let pp_model_private ppf {Program.prepare_data; _} = +let pp_model_private ppf { Program.prepare_data; _ } = let data_decls = List.concat_map ~f:top_level_decls prepare_data |> List.filter_map ~f:ident in pf ppf "%a" (list ~sep:cut pp_decl) data_decls +;; (** Print the signature and blocks of the model class methods. @param ppf A pretty printer @@ -360,71 +401,94 @@ let pp_model_private ppf {Program.prepare_data; _} = @param cv_attr Optional parameter to add method attributes. @param ppbody (?A pretty printer of the method's body) *) -let pp_method ppf rt name params intro ?(outro = []) ?(cv_attr = ["const"]) - ppbody = - pf ppf "@[inline %s %s(@[@,%a@]) %a " rt name - (list ~sep:comma string) params (list ~sep:cut string) cv_attr ; - pf ppf "{@,%a" (list ~sep:cut string) intro ; - pf ppf "@ " ; - ppbody ppf ; - if not (List.is_empty outro) then pf ppf "@ %a" (list ~sep:cut string) outro ; +let pp_method ppf rt name params intro ?(outro = []) ?(cv_attr = [ "const" ]) ppbody = + pf + ppf + "@[inline %s %s(@[@,%a@]) %a " + rt + name + (list ~sep:comma string) + params + (list ~sep:cut string) + cv_attr; + pf ppf "{@,%a" (list ~sep:cut string) intro; + pf ppf "@ "; + ppbody ppf; + if not (List.is_empty outro) then pf ppf "@ %a" (list ~sep:cut string) outro; pf ppf "@,} // %s() @,@]" name +;; (** Print the `get_param_names` method of the model class @param ppf A pretty printer. *) -let pp_get_param_names ppf {Program.output_vars; _} = +let pp_get_param_names ppf { Program.output_vars; _ } = let add_param = fmt "names__.emplace_back(%S);" in - pp_method ppf "void" "get_param_names" ["std::vector& names__"] - [] (fun ppf -> - pf ppf "names__.clear();@ " ; - (list ~sep:cut add_param) ppf (List.map ~f:fst output_vars) ) + pp_method + ppf + "void" + "get_param_names" + [ "std::vector& names__" ] + [] + (fun ppf -> + pf ppf "names__.clear();@ "; + (list ~sep:cut add_param) ppf (List.map ~f:fst output_vars)) +;; (** Print the `get_dims` method of the model class. *) -let pp_get_dims ppf {Program.output_vars; _} = - let pp_cast ppf cast_dims = - pf ppf "static_cast(%a)@," pp_expr cast_dims - in +let pp_get_dims ppf { Program.output_vars; _ } = + let pp_cast ppf cast_dims = pf ppf "static_cast(%a)@," pp_expr cast_dims in let pp_pack ppf inner_dims = - pf ppf "std::vector{@[@,%a@]}" (list ~sep:comma pp_cast) - inner_dims - in - let pp_add_pack ppf dims = - pf ppf "dimss__.emplace_back(%a);@," pp_pack dims + pf ppf "std::vector{@[@,%a@]}" (list ~sep:comma pp_cast) inner_dims in + let pp_add_pack ppf dims = pf ppf "dimss__.emplace_back(%a);@," pp_pack dims in let pp_output_var ppf = (list ~sep:cut pp_add_pack) ppf List.( - map ~f:SizedType.get_dims - (map - ~f:(fun (_, {Program.out_constrained_st= st; _}) -> st) - output_vars)) - in - let params = ["std::vector>& dimss__"] in - let cv_attr = ["const"; "final"] in - pp_method ppf "void" "get_dims" params ["dimss__.clear();"] + map + ~f:SizedType.get_dims + (map ~f:(fun (_, { Program.out_constrained_st = st; _ }) -> st) output_vars)) + in + let params = [ "std::vector>& dimss__" ] in + let cv_attr = [ "const"; "final" ] in + pp_method + ppf + "void" + "get_dims" + params + [ "dimss__.clear();" ] (fun ppf -> pp_output_var ppf) ~cv_attr - -let pp_method_b ppf rt name params intro ?(outro = []) ?(cv_attr = ["const"]) - body = - pp_method ppf rt name params intro +;; + +let pp_method_b ppf rt name params intro ?(outro = []) ?(cv_attr = [ "const" ]) body = + pp_method + ppf + rt + name + params + intro (fun ppf -> pp_located_error_b ppf body) - ~outro ~cv_attr + ~outro + ~cv_attr +;; (** Print the write_array method of the model class *) -let pp_write_array ppf {Program.prog_name; generate_quantities; _} = - pf ppf "template @ " ; +let pp_write_array ppf { Program.prog_name; generate_quantities; _ } = + pf ppf "template @ "; let params = - [ "RNG& base_rng__"; "std::vector& params_r__" - ; "std::vector& params_i__"; "std::vector& vars__" + [ "RNG& base_rng__" + ; "std::vector& params_r__" + ; "std::vector& params_i__" + ; "std::vector& vars__" ; "bool emit_transformed_parameters__ = true" ; "bool emit_generated_quantities__ = true" - ; "std::ostream* pstream__ = nullptr" ] + ; "std::ostream* pstream__ = nullptr" + ] in let intro = - [ "using local_scalar_t__ = double;"; "vars__.resize(0);" + [ "using local_scalar_t__ = double;" + ; "vars__.resize(0);" ; "stan::io::reader in__(params_r__, params_i__);" ; strf "%a" pp_function__ (prog_name, "write_array") ; strf "%a" pp_unused "function__" @@ -432,9 +496,11 @@ let pp_write_array ppf {Program.prog_name; generate_quantities; _} = ; "(void) lp__; // dummy to suppress unused var warning" ; "stan::math::accumulator lp_accum__;" ; "local_scalar_t__ DUMMY_VAR__(std::numeric_limits::quiet_NaN());" - ; strf "%a" pp_unused "DUMMY_VAR__" ] + ; strf "%a" pp_unused "DUMMY_VAR__" + ] in pp_method_b ppf "void" "write_array" params intro generate_quantities +;; (** Prints the for loop for `constrained_param_names` and `unconstrained_param_names` @@ -446,42 +512,46 @@ let pp_write_array ppf {Program.prog_name; generate_quantities; _} = let rec pp_for_loop_iteratee ?(index_ids = []) ppf (iteratee, dims, pp_body) = let iter d pp_body = let loopvar, gensym_exit = Common.Gensym.enter () in - pp_for_loop ppf + pp_for_loop + ppf ( loopvar , Expr.Helpers.loop_bottom , d , pp_block - , (pp_body, (iteratee, loopvar :: index_ids)) ) ; + , (pp_body, (iteratee, loopvar :: index_ids)) ); gensym_exit () in match dims with | [] -> pp_body ppf (iteratee, index_ids) | dim :: dims -> - iter dim (fun ppf (i, idcs) -> - pf ppf "%a" pp_block - (pp_for_loop_iteratee ~index_ids:idcs, (i, dims, pp_body)) ) + iter dim (fun ppf (i, idcs) -> + pf ppf "%a" pp_block (pp_for_loop_iteratee ~index_ids:idcs, (i, dims, pp_body))) +;; (** Print the `constrained_param_names` method of the model class. *) -let pp_constrained_param_names ppf {Program.output_vars; _} = +let pp_constrained_param_names ppf { Program.output_vars; _ } = let params = [ "std::vector& param_names__" ; "bool emit_transformed_parameters__ = true" - ; "bool emit_generated_quantities__ = true" ] + ; "bool emit_generated_quantities__ = true" + ] in let paramvars, tparamvars, gqvars = List.partition3_map ~f:(function - | id, {Program.out_block= Parameters; out_constrained_st= st; _} -> - `Fst (id, st) - | id, {out_block= TransformedParameters; out_constrained_st= st; _} -> - `Snd (id, st) - | id, {out_block= GeneratedQuantities; out_constrained_st= st; _} -> - `Trd (id, st)) + | id, { Program.out_block = Parameters; out_constrained_st = st; _ } -> + `Fst (id, st) + | id, { out_block = TransformedParameters; out_constrained_st = st; _ } -> + `Snd (id, st) + | id, { out_block = GeneratedQuantities; out_constrained_st = st; _ } -> + `Trd (id, st)) output_vars in let emit_name ppf (name, idcs) = let to_string = fmt "std::to_string(%s)" in - pf ppf "param_names__.emplace_back(std::string() + %a);" + pf + ppf + "param_names__.emplace_back(std::string() + %a);" (list ~sep:(fun ppf () -> pf ppf " + '.' + ") string) (strf "%S" name :: List.map ~f:(strf "%a" to_string) idcs) in @@ -489,15 +559,27 @@ let pp_constrained_param_names ppf {Program.output_vars; _} = let dims = List.rev (SizedType.get_dims st) in pp_for_loop_iteratee ppf (decl_id, dims, emit_name) in - let cv_attr = ["const"; "final"] in - pp_method ppf "void" "constrained_param_names" params [] + let cv_attr = [ "const"; "final" ] in + pp_method + ppf + "void" + "constrained_param_names" + params + [] (fun ppf -> - (list ~sep:cut pp_param_names) ppf paramvars ; - pf ppf "@,if (emit_transformed_parameters__) %a@," pp_block - (list ~sep:cut pp_param_names, tparamvars) ; - pf ppf "@,if (emit_generated_quantities__) %a@," pp_block - (list ~sep:cut pp_param_names, gqvars) ) + (list ~sep:cut pp_param_names) ppf paramvars; + pf + ppf + "@,if (emit_transformed_parameters__) %a@," + pp_block + (list ~sep:cut pp_param_names, tparamvars); + pf + ppf + "@,if (emit_generated_quantities__) %a@," + pp_block + (list ~sep:cut pp_param_names, gqvars)) ~cv_attr +;; (* Print the `unconstrained_param_names` method of the model class. This is just a copy of constrained, I need to figure out which one is wrong @@ -516,27 +598,29 @@ let pp_constrained_param_names ppf {Program.output_vars; _} = change size. The ordered types and constrained types don't change sizes either. *) -let pp_unconstrained_param_names ppf {Program.output_vars; _} = +let pp_unconstrained_param_names ppf { Program.output_vars; _ } = let params = [ "std::vector& param_names__" ; "bool emit_transformed_parameters__ = true" - ; "bool emit_generated_quantities__ = true" ] + ; "bool emit_generated_quantities__ = true" + ] in let paramvars, tparamvars, gqvars = List.partition3_map ~f:(function - | id, {Program.out_block= Parameters; out_unconstrained_st= st; _} -> - `Fst (id, st) - | id, {out_block= TransformedParameters; out_unconstrained_st= st; _} - -> - `Snd (id, st) - | id, {out_block= GeneratedQuantities; out_unconstrained_st= st; _} -> - `Trd (id, st)) + | id, { Program.out_block = Parameters; out_unconstrained_st = st; _ } -> + `Fst (id, st) + | id, { out_block = TransformedParameters; out_unconstrained_st = st; _ } -> + `Snd (id, st) + | id, { out_block = GeneratedQuantities; out_unconstrained_st = st; _ } -> + `Trd (id, st)) output_vars in let emit_name ppf (name, idcs) = let to_string = fmt "std::to_string(%s)" in - pf ppf "param_names__.emplace_back(std::string() + %a);" + pf + ppf + "param_names__.emplace_back(std::string() + %a);" (list ~sep:(fun ppf () -> pf ppf " + '.' + ") string) (strf "%S" name :: List.map ~f:(strf "%a" to_string) idcs) in @@ -544,48 +628,70 @@ let pp_unconstrained_param_names ppf {Program.output_vars; _} = let dims = List.rev (SizedType.get_dims st) in pp_for_loop_iteratee ppf (decl_id, dims, emit_name) in - let cv_attr = ["const"; "final"] in - pp_method ppf "void" "unconstrained_param_names" params [] + let cv_attr = [ "const"; "final" ] in + pp_method + ppf + "void" + "unconstrained_param_names" + params + [] (fun ppf -> - (list ~sep:cut pp_param_names) ppf paramvars ; - pf ppf "@,if (emit_transformed_parameters__) %a@," pp_block - (list ~sep:cut pp_param_names, tparamvars) ; - pf ppf "@,if (emit_generated_quantities__) %a@," pp_block - (list ~sep:cut pp_param_names, gqvars) ) + (list ~sep:cut pp_param_names) ppf paramvars; + pf + ppf + "@,if (emit_transformed_parameters__) %a@," + pp_block + (list ~sep:cut pp_param_names, tparamvars); + pf + ppf + "@,if (emit_generated_quantities__) %a@," + pp_block + (list ~sep:cut pp_param_names, gqvars)) ~cv_attr +;; (** Print the `transform_inits` method of the model class *) -let pp_transform_inits ppf {Program.transform_inits; _} = +let pp_transform_inits ppf { Program.transform_inits; _ } = let params = - [ "const stan::io::var_context& context__"; "std::vector& params_i__" - ; "std::vector& vars__"; "std::ostream* pstream__" ] + [ "const stan::io::var_context& context__" + ; "std::vector& params_i__" + ; "std::vector& vars__" + ; "std::ostream* pstream__" + ] in let intro = - [ "using local_scalar_t__ = double;"; "vars__.clear();" - ; "vars__.reserve(num_params_r__);" ] + [ "using local_scalar_t__ = double;" + ; "vars__.clear();" + ; "vars__.reserve(num_params_r__);" + ] in - let cv_attr = ["const"; "final"] in - pp_method_b ppf "void" "transform_inits" params intro transform_inits - ~cv_attr + let cv_attr = [ "const"; "final" ] in + pp_method_b ppf "void" "transform_inits" params intro transform_inits ~cv_attr +;; (** Print the `log_prob` method of the model class *) -let pp_log_prob ppf Program.({prog_name; log_prob; _}) = - pf ppf "template @ " ; +let pp_log_prob ppf Program.{ prog_name; log_prob; _ } = + pf ppf "template @ "; let params = - [ "std::vector& params_r__"; "std::vector& params_i__" - ; "std::ostream* pstream__ = nullptr" ] + [ "std::vector& params_r__" + ; "std::vector& params_i__" + ; "std::ostream* pstream__ = nullptr" + ] in let intro = - [ "using local_scalar_t__ = T__;"; "T__ lp__(0.0);" + [ "using local_scalar_t__ = T__;" + ; "T__ lp__(0.0);" ; "stan::math::accumulator lp_accum__;" ; strf "%a" pp_function__ (prog_name, "log_prob") ; "stan::io::reader in__(params_r__, params_i__);" ; "local_scalar_t__ DUMMY_VAR__(std::numeric_limits::quiet_NaN());" - ; strf "%a" pp_unused "DUMMY_VAR__" ] + ; strf "%a" pp_unused "DUMMY_VAR__" + ] in - let outro = ["lp_accum__.add(lp__);"; "return lp_accum__.sum();"] in - let cv_attr = ["const"] in + let outro = [ "lp_accum__.add(lp__);"; "return lp_accum__.sum();" ] in + let cv_attr = [ "const" ] in pp_method_b ppf "T__" "log_prob" params intro log_prob ~outro ~cv_attr +;; (** Print the body of the constrained and unconstrained sizedtype methods in the model class @@ -594,31 +700,35 @@ let pp_log_prob ppf Program.({prog_name; log_prob; _}) = @param outvars The parameters to gather the sizes for. *) let pp_outvar_metadata ppf (method_name, outvars) = - let intro = ["stringstream s__;"] in - let outro = ["return s__.str();"] in + let intro = [ "stringstream s__;" ] in + let outro = [ "return s__.str();" ] in let json_str = Cpp_Json.out_var_interpolated_json_str outvars in let ppbody ppf = pf ppf "s__ << %s;" json_str in pp_method ppf "std::string" method_name [] intro ~outro ppbody +;; (** Print the `get_unconstrained_sizedtypes` method of the model class *) -let pp_unconstrained_types ppf {Program.output_vars; _} = - let grab_unconstrained (name, {Program.out_unconstrained_st; out_block; _}) = - (name, out_unconstrained_st, out_block) +let pp_unconstrained_types ppf { Program.output_vars; _ } = + let grab_unconstrained (name, { Program.out_unconstrained_st; out_block; _ }) = + name, out_unconstrained_st, out_block in let outvars = List.map ~f:grab_unconstrained output_vars in pp_outvar_metadata ppf ("get_unconstrained_sizedtypes", outvars) +;; (** Print the `get_constrained_sizedtypes` method of the model class *) -let pp_constrained_types ppf {Program.output_vars; _} = - let grab_constrained (name, {Program.out_constrained_st; out_block; _}) = - (name, out_constrained_st, out_block) +let pp_constrained_types ppf { Program.output_vars; _ } = + let grab_constrained (name, { Program.out_constrained_st; out_block; _ }) = + name, out_constrained_st, out_block in let outvars = List.map ~f:grab_constrained output_vars in pp_outvar_metadata ppf ("get_constrained_sizedtypes", outvars) +;; (** Print the generic method overloads needed in the model class. *) let pp_overloads ppf () = - pf ppf + pf + ppf {| // Begin method overload boilerplate template @@ -662,32 +772,34 @@ let pp_overloads ppf () = params_r(i) = params_r_vec[i]; } |} +;; (** Print the public parts of the model class *) let pp_model_public ppf p = - pf ppf "@ %a" pp_ctor p ; - pf ppf "@ %a" pp_log_prob p ; - pf ppf "@ %a" pp_write_array p ; - pf ppf "@ %a" pp_transform_inits p ; + pf ppf "@ %a" pp_ctor p; + pf ppf "@ %a" pp_log_prob p; + pf ppf "@ %a" pp_write_array p; + pf ppf "@ %a" pp_transform_inits p; (* Begin metadata methods *) - pf ppf "@ %a" pp_get_param_names p ; + pf ppf "@ %a" pp_get_param_names p; (* Post-data metadata methods *) - pf ppf "@ %a" pp_get_dims p ; - pf ppf "@ %a" pp_constrained_param_names p ; - pf ppf "@ %a" pp_unconstrained_param_names p ; - pf ppf "@ %a" pp_constrained_types p ; - pf ppf "@ %a" pp_unconstrained_types p ; + pf ppf "@ %a" pp_get_dims p; + pf ppf "@ %a" pp_constrained_param_names p; + pf ppf "@ %a" pp_unconstrained_param_names p; + pf ppf "@ %a" pp_constrained_types p; + pf ppf "@ %a" pp_unconstrained_types p; (* Boilerplate *) pf ppf "@ %a" pp_overloads () +;; (** Print the full model class. *) -let pp_model ppf ({Program.prog_name; _} as p) = - pf ppf "class %s final : public model_base_crtp<%s> {" prog_name prog_name ; - pf ppf "@ @[@ private:@ @[ %a@]@ " pp_model_private p ; - pf ppf "@ public:@ @[ ~%s() final { }" p.prog_name ; - pf ppf "@ @ std::string model_name() const final { return \"%s\"; }" - prog_name ; - pf ppf +let pp_model ppf ({ Program.prog_name; _ } as p) = + pf ppf "class %s final : public model_base_crtp<%s> {" prog_name prog_name; + pf ppf "@ @[@ private:@ @[ %a@]@ " pp_model_private p; + pf ppf "@ public:@ @[ ~%s() final { }" p.prog_name; + pf ppf "@ @ std::string model_name() const final { return \"%s\"; }" prog_name; + pf + ppf {| std::vector model_compile_info() const { @@ -697,8 +809,10 @@ let pp_model ppf ({Program.prog_name; _} as p) = return stanc_info; } |} - "%%NAME%%3 %%VERSION%%" stanc_args_to_print ; + "%%NAME%%3 %%VERSION%%" + stanc_args_to_print; pf ppf "@ %a@]@]@ };" pp_model_public p +;; (** The C++ aliases needed for the model class*) let usings = @@ -721,6 +835,7 @@ using stan::model::index_multi; using stan::model::index_omni; using stan::model::nil_index_list; using namespace stan::math; |} +;; (** Functions needed in the model class not defined yet in stan math. FIXME: Move these to the Stan repo when these repos are joined. @@ -759,33 +874,37 @@ inline void validate_unit_vector_index(const char* var_name, const char* expr, } } |} +;; (** Create the model's namespace. *) -let namespace Program.({prog_name; _}) = prog_name ^ "_namespace" +let namespace Program.{ prog_name; _ } = prog_name ^ "_namespace" (** Find and register functiors used for map_rect. *) let pp_register_map_rect_functors ppf p = let pp_register_functor ppf (i, f) = pf ppf "STAN_REGISTER_MAP_RECT(%d, %s::%s)" i (namespace p) f in - pf ppf "@ %a" + pf + ppf + "@ %a" (list ~sep:cut pp_register_functor) (List.sort ~compare (Hashtbl.to_alist map_rect_calls)) +;; let fun_used_in_reduce_sum p = - let rec find_functors_expr accum Expr.Fixed.({pattern; _}) = - String.Set.union accum - ( match pattern with - | FunApp (StanLib, x, {pattern= Var f; _} :: _) - when Stan_math_signatures.is_reduce_sum_fn x -> - String.Set.of_list [f] - | x -> Expr.Fixed.Pattern.fold find_functors_expr accum x ) + let rec find_functors_expr accum Expr.Fixed.{ pattern; _ } = + String.Set.union + accum + (match pattern with + | FunApp (StanLib, x, { pattern = Var f; _ } :: _) + when Stan_math_signatures.is_reduce_sum_fn x -> String.Set.of_list [ f ] + | x -> Expr.Fixed.Pattern.fold find_functors_expr accum x) in let rec find_functors_stmt accum stmt = - Stmt.Fixed.( - Pattern.fold find_functors_expr find_functors_stmt accum stmt.pattern) + Stmt.Fixed.(Pattern.fold find_functors_expr find_functors_stmt accum stmt.pattern) in Program.fold find_functors_expr find_functors_stmt String.Set.empty p +;; (** Print the full C++ for the stan program. *) let pp_prog ppf (p : Program.Typed.t) = @@ -799,13 +918,24 @@ let pp_prog ppf (p : Program.Typed.t) = ~f:(fun x -> "struct " ^ x ^ reduce_sum_functor_suffix ^ ";") (fun_used_in_reduce_sum p) in - pf ppf "@[@ %s@ %s@ namespace %s {@ %s@ %s@ %a@ %s@ %a@ %a@ }@ @]" version - includes (namespace p) custom_functions usings Locations.pp_globals s + pf + ppf + "@[@ %s@ %s@ namespace %s {@ %s@ %s@ %a@ %s@ %a@ %a@ }@ @]" + version + includes + (namespace p) + custom_functions + usings + Locations.pp_globals + s (String.concat ~sep:"\n" (String.Set.elements reduce_sum_struct_decl)) (list ~sep:cut pp_fun_def_with_rs_list) - p.functions_block pp_model p ; - pf ppf "@,using stan_model = %s_namespace::%s;@," p.prog_name p.prog_name ; - pf ppf + p.functions_block + pp_model + p; + pf ppf "@,using stan_model = %s_namespace::%s;@," p.prog_name p.prog_name; + pf + ppf {| #ifndef USING_R @@ -819,5 +949,6 @@ stan::model::model_base& new_model( } #endif -|} ; +|}; pf ppf "@[%a@]" pp_register_map_rect_functors p +;; diff --git a/src/stan_math_backend/Statement_gen.ml b/src/stan_math_backend/Statement_gen.ml index dee4c2f418..dc9d270616 100644 --- a/src/stan_math_backend/Statement_gen.ml +++ b/src/stan_math_backend/Statement_gen.ml @@ -10,6 +10,7 @@ let rec contains_eigen = function | UnsizedType.UArray t -> contains_eigen t | UMatrix | URowVector | UVector -> true | UInt | UReal | UMathLibraryFunction | UFun _ -> false +;; let pp_set_size ppf (decl_id, st, adtype) = (* TODO: generate optimal adtypes for expressions and declarations *) @@ -29,51 +30,74 @@ let pp_set_size ppf (decl_id, st, adtype) = | SMatrix (d1, d2) -> pf ppf "%a(%a, %a)" pp_st st pp_expr d1 pp_expr d2 | SArray (t, d) -> pf ppf "%a(%a, %a)" pp_st st pp_expr d pp_size_ctor t in - pf ppf "@[%s = %a;@]@," decl_id pp_size_ctor st ; - if contains_eigen (SizedType.to_unsized st) then - pf ppf "@[stan::math::fill(%s, %s);@]@," decl_id real_nan + pf ppf "@[%s = %a;@]@," decl_id pp_size_ctor st; + if contains_eigen (SizedType.to_unsized st) + then pf ppf "@[stan::math::fill(%s, %s);@]@," decl_id real_nan +;; let%expect_test "set size mat array" = let int = Expr.Helpers.int in - strf "@[%a@]" pp_set_size + strf + "@[%a@]" + pp_set_size ("d", SArray (SArray (SMatrix (int 2, int 3), int 4), int 5), DataOnly) - |> print_endline ; + |> print_endline; [%expect {| d = std::vector>>(5, std::vector>(4, Eigen::Matrix(2, 3))); stan::math::fill(d, std::numeric_limits::quiet_NaN()); |}] +;; (** [pp_for_loop ppf (loopvar, lower, upper, pp_body, body)] tries to pretty print a for-loop from lower to upper given some loopvar.*) let pp_for_loop ppf (loopvar, lower, upper, pp_body, body) = - pf ppf "@[for (@[int %s = %a;@ %s <= %a;@ ++%s@])" loopvar pp_expr - lower loopvar pp_expr upper loopvar ; + pf + ppf + "@[for (@[int %s = %a;@ %s <= %a;@ ++%s@])" + loopvar + pp_expr + lower + loopvar + pp_expr + upper + loopvar; pf ppf " %a@]" pp_body body +;; let rec integer_el_type = function | SizedType.SReal | SVector _ | SMatrix _ | SRowVector _ -> false | SInt -> true | SArray (st, _) -> integer_el_type st +;; let pp_decl ppf (vident, ut, adtype) = let pp_type = - if Transform_Mir.is_opencl_var vident then fun ppf _ -> + if Transform_Mir.is_opencl_var vident + then + fun ppf _ -> match ut with | UnsizedType.UInt | UArray UInt -> pf ppf "matrix_cl" | _ -> pf ppf "matrix_cl" else pp_unsizedtype_local in pf ppf "%a %s;" pp_type (adtype, ut) vident +;; let pp_sized_decl ppf (vident, st, adtype) = - pf ppf "%a@,%a" pp_decl + pf + ppf + "%a@,%a" + pp_decl (vident, SizedType.to_unsized st, adtype) - pp_set_size (vident, st, adtype) + pp_set_size + (vident, st, adtype) +;; let pp_possibly_sized_decl ppf (vident, pst, adtype) = match pst with | Type.Sized st -> pp_sized_decl ppf (vident, st, adtype) | Unsized ut -> pp_decl ppf (vident, ut, adtype) +;; let math_fn_translations = function | Internal_fun.FnLength -> Some ("length", []) @@ -81,127 +105,139 @@ let math_fn_translations = function | FnValidateSizeSimplex -> Some ("validate_positive_index", []) | FnValidateSizeUnitVector -> Some ("validate_unit_vector_index", []) | _ -> None +;; let trans_math_fn fname = Option.( - value ~default:(fname, []) + value + ~default:(fname, []) (bind (Internal_fun.of_string_opt fname) ~f:math_fn_translations)) +;; let pp_bool_expr ppf expr = match Expr.Typed.type_of expr with - | UReal -> pp_call ppf ("as_bool", pp_expr, [expr]) + | UReal -> pp_call ppf ("as_bool", pp_expr, [ expr ]) | _ -> pp_expr ppf expr +;; -let rec pp_statement (ppf : Format.formatter) - (Stmt.Fixed.({pattern; meta}) as stmt) = +let rec pp_statement (ppf : Format.formatter) (Stmt.Fixed.{ pattern; meta } as stmt) = (* ({stmt; smeta} : (mtype_loc_ad, 'a) stmt_with) = *) let pp_stmt_list = list ~sep:cut pp_statement in - ( match pattern with + (match pattern with | Block _ | SList _ | Decl _ | Skip | Break | Continue -> () - | _ -> Locations.pp_smeta ppf meta ) ; + | _ -> Locations.pp_smeta ppf meta); match pattern with | Assignment - ((vident, _, []), ({meta= Expr.Typed.Meta.({type_= UInt; _}); _} as rhs)) - |Assignment ((vident, _, []), ({meta= {type_= UReal; _}; _} as rhs)) -> - pf ppf "@[%s = %a;@]" vident pp_expr rhs - | Assignment ((assignee, UInt, idcs), rhs) - |Assignment ((assignee, UReal, idcs), rhs) + ((vident, _, []), ({ meta = Expr.Typed.Meta.{ type_ = UInt; _ }; _ } as rhs)) + | Assignment ((vident, _, []), ({ meta = { type_ = UReal; _ }; _ } as rhs)) -> + pf ppf "@[%s = %a;@]" vident pp_expr rhs + | (Assignment ((assignee, UInt, idcs), rhs) | Assignment ((assignee, UReal, idcs), rhs)) when List.for_all ~f:is_single_index idcs -> - pf ppf "@[%a = %a;@]" pp_indexed_simple (assignee, idcs) pp_expr - rhs + pf ppf "@[%a = %a;@]" pp_indexed_simple (assignee, idcs) pp_expr rhs | Assignment ((assignee, _, idcs), rhs) -> - (* XXX I think in general we don't need to do a deepcopy if e is nested + (* XXX I think in general we don't need to do a deepcopy if e is nested inside some function call - the function should get its own copy (in all cases???) *) - let rec maybe_deep_copy e = - let recurse (e : 'a Expr.Fixed.t) = - { e with - Expr.Fixed.pattern= - Expr.Fixed.Pattern.map maybe_deep_copy e.pattern } - in - match e.pattern with - | _ when UnsizedType.is_scalar_type (Expr.Typed.type_of e) -> e - | FunApp (CompilerInternal, _, _) -> e - | (Indexed ({Expr.Fixed.pattern= Var v; _}, _) | Var v) - when v = assignee -> - { e with - Expr.Fixed.pattern= - FunApp (CompilerInternal, "stan::model::deep_copy", [e]) } - | _ -> recurse e + let rec maybe_deep_copy e = + let recurse (e : 'a Expr.Fixed.t) = + { e with Expr.Fixed.pattern = Expr.Fixed.Pattern.map maybe_deep_copy e.pattern } in - let rhs = - match rhs.pattern with - | FunApp (CompilerInternal, f, _) - when f = Internal_fun.to_string FnConstrain - || f = Internal_fun.to_string FnUnconstrain -> - rhs - | _ -> maybe_deep_copy rhs - in - pf ppf "@[assign(@,%s,@ %a,@ %a,@ %S@]);" assignee pp_indexes idcs - pp_expr rhs - (strf "assigning variable %s" - assignee - (* (list ~sep:comma (Pretty.pp_index Pretty.pp_expr_typed_located)) idcs *)) + match e.pattern with + | _ when UnsizedType.is_scalar_type (Expr.Typed.type_of e) -> e + | FunApp (CompilerInternal, _, _) -> e + | (Indexed ({ Expr.Fixed.pattern = Var v; _ }, _) | Var v) when v = assignee -> + { e with + Expr.Fixed.pattern = FunApp (CompilerInternal, "stan::model::deep_copy", [ e ]) + } + | _ -> recurse e + in + let rhs = + match rhs.pattern with + | FunApp (CompilerInternal, f, _) + when f = Internal_fun.to_string FnConstrain + || f = Internal_fun.to_string FnUnconstrain -> rhs + | _ -> maybe_deep_copy rhs + in + pf + ppf + "@[assign(@,%s,@ %a,@ %a,@ %S@]);" + assignee + pp_indexes + idcs + pp_expr + rhs + (strf + "assigning variable %s" + assignee + (* (list ~sep:comma (Pretty.pp_index Pretty.pp_expr_typed_located)) idcs *)) | TargetPE e -> pf ppf "@[lp_accum__.add(@,%a@]);" pp_expr e - | NRFunApp (CompilerInternal, fname, args) - when fname = Internal_fun.to_string FnPrint -> - let pp_arg ppf a = pf ppf "stan_print(pstream__, %a);" pp_expr a in - let args = args @ [Expr.Helpers.str "\n"] in - pf ppf "if (pstream__) %a" pp_block (list ~sep:cut pp_arg, args) - | NRFunApp (CompilerInternal, fname, args) - when fname = Internal_fun.to_string FnReject -> - let err_strm = "errmsg_stream__" in - let add_to_string ppf e = pf ppf "%s << %a;" err_strm pp_expr e in - pf ppf "std::stringstream %s;@," err_strm ; - pf ppf "%a@," (list ~sep:cut add_to_string) args ; - pf ppf "throw std::domain_error(%s.str());" err_strm - | NRFunApp - (CompilerInternal, fname, {pattern= Lit (Str, check_name); _} :: args) + | NRFunApp (CompilerInternal, fname, args) when fname = Internal_fun.to_string FnPrint + -> + let pp_arg ppf a = pf ppf "stan_print(pstream__, %a);" pp_expr a in + let args = args @ [ Expr.Helpers.str "\n" ] in + pf ppf "if (pstream__) %a" pp_block (list ~sep:cut pp_arg, args) + | NRFunApp (CompilerInternal, fname, args) when fname = Internal_fun.to_string FnReject + -> + let err_strm = "errmsg_stream__" in + let add_to_string ppf e = pf ppf "%s << %a;" err_strm pp_expr e in + pf ppf "std::stringstream %s;@," err_strm; + pf ppf "%a@," (list ~sep:cut add_to_string) args; + pf ppf "throw std::domain_error(%s.str());" err_strm + | NRFunApp (CompilerInternal, fname, { pattern = Lit (Str, check_name); _ } :: args) when fname = Internal_fun.to_string FnCheck -> - let args = - {Expr.Fixed.pattern= Var "function__"; meta= Expr.Typed.Meta.empty} - :: args - in - pp_statement ppf - { pattern= NRFunApp (CompilerInternal, "check_" ^ check_name, args) - ; meta= stmt.meta } - | NRFunApp (CompilerInternal, fname, [var]) + let args = + { Expr.Fixed.pattern = Var "function__"; meta = Expr.Typed.Meta.empty } :: args + in + pp_statement + ppf + { pattern = NRFunApp (CompilerInternal, "check_" ^ check_name, args) + ; meta = stmt.meta + } + | NRFunApp (CompilerInternal, fname, [ var ]) when fname = Internal_fun.to_string FnWriteParam -> - pf ppf "@[vars__.emplace_back(@,%a);@]" pp_expr var + pf ppf "@[vars__.emplace_back(@,%a);@]" pp_expr var | NRFunApp (CompilerInternal, fname, args) -> - let fname, extra_args = trans_math_fn fname in - pf ppf "%s(@[%a@]);" fname (list ~sep:comma pp_expr) - (extra_args @ args) + let fname, extra_args = trans_math_fn fname in + pf ppf "%s(@[%a@]);" fname (list ~sep:comma pp_expr) (extra_args @ args) | NRFunApp (StanLib, fname, args) -> - pf ppf "%s(@[%a@]);" fname (list ~sep:comma pp_expr) args - | NRFunApp (UserDefined, fname, args) -> - pf ppf "%a;" pp_user_defined_fun (fname, args) + pf ppf "%s(@[%a@]);" fname (list ~sep:comma pp_expr) args + | NRFunApp (UserDefined, fname, args) -> pf ppf "%a;" pp_user_defined_fun (fname, args) | Break -> string ppf "break;" | Continue -> string ppf "continue;" | Return e -> pf ppf "@[return %a;@]" (option pp_expr) e | Skip -> string ppf ";" | IfElse (cond, ifbranch, elsebranch) -> - let pp_else ppf x = pf ppf "else %a" pp_statement x in - pf ppf "if (@[%a@]) %a %a" pp_bool_expr cond pp_block_s ifbranch - (option pp_else) elsebranch + let pp_else ppf x = pf ppf "else %a" pp_statement x in + pf + ppf + "if (@[%a@]) %a %a" + pp_bool_expr + cond + pp_block_s + ifbranch + (option pp_else) + elsebranch | While (cond, body) -> - pf ppf "while (@[%a@]) %a" pp_bool_expr cond pp_block_s body + pf ppf "while (@[%a@]) %a" pp_bool_expr cond pp_block_s body | For - { body= - { pattern= - Assignment (_, {pattern= FunApp (CompilerInternal, f, _); _}); _ - } as body; _ } + { body = + { pattern = Assignment (_, { pattern = FunApp (CompilerInternal, f, _); _ }) + ; _ + } as body + ; _ + } when Internal_fun.of_string_opt f = Some FnReadParam -> - pp_statement ppf body - (* Skip For loop part, just emit body due to the way FnReadParam emits *) - | For {loopvar; lower; upper; body} -> - pp_for_loop ppf (loopvar, lower, upper, pp_statement, body) + pp_statement ppf body + (* Skip For loop part, just emit body due to the way FnReadParam emits *) + | For { loopvar; lower; upper; body } -> + pp_for_loop ppf (loopvar, lower, upper, pp_statement, body) | Block ls -> pp_block ppf (pp_stmt_list, ls) | SList ls -> pp_stmt_list ppf ls - | Decl {decl_adtype; decl_id; decl_type} -> - pp_possibly_sized_decl ppf (decl_id, decl_type, decl_adtype) + | Decl { decl_adtype; decl_id; decl_type } -> + pp_possibly_sized_decl ppf (decl_id, decl_type, decl_adtype) and pp_block_s ppf body = match body.pattern with | Block ls -> pp_block ppf (list ~sep:cut pp_statement, ls) | _ -> pp_block ppf (pp_statement, body) +;; diff --git a/src/stan_math_backend/Transform_Mir.ml b/src/stan_math_backend/Transform_Mir.ml index 036bbf62db..d193d62efb 100644 --- a/src/stan_math_backend/Transform_Mir.ml +++ b/src/stan_math_backend/Transform_Mir.ml @@ -6,35 +6,31 @@ let use_opencl = ref false let opencl_triggers = String.Map.of_alist_exn [ ( "normal_id_glm_lpdf" - , ( [0; 1] + , ( [ 0; 1 ] , [ (* Array of conditions under which to move to OpenCL *) - ([1], (* Argument 1 is data *) - [(1, UnsizedType.UMatrix)]) - (* Argument 1 is a matrix *) - ] ) ) - ; ( "bernoulli_logit_glm_lpmf" - , ([0; 1], [([1], [(1, UnsizedType.UMatrix)])]) ) - ; ( "categorical_logit_glm_lpmf" - , ([0; 1], [([1], [(1, UnsizedType.UMatrix)])]) ) - ; ( "neg_binomial_2_log_glm_lpmf" - , ([0; 1], [([1], [(1, UnsizedType.UMatrix)])]) ) - ; ( "ordered_logistic_glm_lpmf" - , ([0; 1], [([1], [(1, UnsizedType.UMatrix)])]) ) - ; ("poisson_log_glm_lpmf", ([0; 1], [([1], [(1, UnsizedType.UMatrix)])])) + [ 1 ], (* Argument 1 is data *) + [ 1, UnsizedType.UMatrix ] + (* Argument 1 is a matrix *) + ] ) ) + ; "bernoulli_logit_glm_lpmf", ([ 0; 1 ], [ [ 1 ], [ 1, UnsizedType.UMatrix ] ]) + ; "categorical_logit_glm_lpmf", ([ 0; 1 ], [ [ 1 ], [ 1, UnsizedType.UMatrix ] ]) + ; "neg_binomial_2_log_glm_lpmf", ([ 0; 1 ], [ [ 1 ], [ 1, UnsizedType.UMatrix ] ]) + ; "ordered_logistic_glm_lpmf", ([ 0; 1 ], [ [ 1 ], [ 1, UnsizedType.UMatrix ] ]) + ; "poisson_log_glm_lpmf", ([ 0; 1 ], [ [ 1 ], [ 1, UnsizedType.UMatrix ] ]) ] +;; let opencl_suffix = "_opencl__" let to_matrix_cl e = - Expr.Fixed.{e with pattern= FunApp (StanLib, "to_matrix_cl", [e])} + Expr.Fixed.{ e with pattern = FunApp (StanLib, "to_matrix_cl", [ e ]) } +;; -let rec switch_expr_to_opencl available_cl_vars (Expr.Fixed.({pattern; _}) as e) - = +let rec switch_expr_to_opencl available_cl_vars (Expr.Fixed.{ pattern; _ } as e) = let is_avail = List.mem available_cl_vars ~equal:( = ) in - let to_cl (Expr.Fixed.({pattern; _}) as e) = + let to_cl (Expr.Fixed.{ pattern; _ } as e) = match pattern with - | Var s when is_avail s -> - Expr.Fixed.{e with pattern= Var (s ^ opencl_suffix)} + | Var s when is_avail s -> Expr.Fixed.{ e with pattern = Var (s ^ opencl_suffix) } | _ -> to_matrix_cl e in let move_cl_args cl_args index arg = @@ -42,8 +38,10 @@ let rec switch_expr_to_opencl available_cl_vars (Expr.Fixed.({pattern; _}) as e) in let check_type args (i, t) = Expr.Typed.type_of (List.nth_exn args i) = t in let check_if_data args ind = - let Expr.Fixed.({pattern; _}) = List.nth_exn args ind in - match pattern with Var s when is_avail s -> true | _ -> false + let Expr.Fixed.{ pattern; _ } = List.nth_exn args ind in + match pattern with + | Var s when is_avail s -> true + | _ -> false in let req_met args (data_arg, type_arg) = List.for_all ~f:(check_if_data args) data_arg @@ -58,19 +56,19 @@ let rec switch_expr_to_opencl available_cl_vars (Expr.Fixed.({pattern; _}) as e) match pattern with | FunApp (StanLib, f, args) when Map.mem opencl_triggers (Utils.stdlib_distribution_name f) -> - let trigger = - Map.find_exn opencl_triggers (Utils.stdlib_distribution_name f) - in - {e with pattern= FunApp (StanLib, f, maybe_map_args args trigger)} + let trigger = Map.find_exn opencl_triggers (Utils.stdlib_distribution_name f) in + { e with pattern = FunApp (StanLib, f, maybe_map_args args trigger) } | x -> - { e with - pattern= - Expr.Fixed.Pattern.map (switch_expr_to_opencl available_cl_vars) x } + { e with + pattern = Expr.Fixed.Pattern.map (switch_expr_to_opencl available_cl_vars) x + } +;; let rec base_type = function | SizedType.SArray (t, _) -> base_type t | SVector _ | SRowVector _ | SMatrix _ -> UnsizedType.UReal | x -> SizedType.to_unsized x +;; let pos = "pos__" @@ -79,64 +77,69 @@ let data_read smeta (decl_id, st) = let scalar = base_type st in let flat_type = UnsizedType.UArray scalar in let decl_var = - { Expr.Fixed.pattern= Var decl_id - ; meta= Expr.Typed.Meta.{loc= smeta; type_= unsized; adlevel= DataOnly} } + { Expr.Fixed.pattern = Var decl_id + ; meta = Expr.Typed.Meta.{ loc = smeta; type_ = unsized; adlevel = DataOnly } + } in - let swrap stmt = {Stmt.Fixed.pattern= stmt; meta= smeta} in - let pos_var = {Expr.Fixed.pattern= Var pos; meta= Expr.Typed.Meta.empty} in + let swrap stmt = { Stmt.Fixed.pattern = stmt; meta = smeta } in + let pos_var = { Expr.Fixed.pattern = Var pos; meta = Expr.Typed.Meta.empty } in let readfnapp var = - Expr.Helpers.internal_funapp FnReadData - [{var with pattern= Lit (Str, decl_id)}] - Expr.Typed.Meta.{var.meta with type_= flat_type} + Expr.Helpers.internal_funapp + FnReadData + [ { var with pattern = Lit (Str, decl_id) } ] + Expr.Typed.Meta.{ var.meta with type_ = flat_type } in match unsized with | UInt | UReal -> - [ Assignment - ( (decl_id, unsized, []) - , { Expr.Fixed.pattern= - Indexed (readfnapp decl_var, [Single Expr.Helpers.loop_bottom]) - ; meta= {decl_var.meta with type_= unsized} } ) - |> swrap ] + [ Assignment + ( (decl_id, unsized, []) + , { Expr.Fixed.pattern = + Indexed (readfnapp decl_var, [ Single Expr.Helpers.loop_bottom ]) + ; meta = { decl_var.meta with type_ = unsized } + } ) + |> swrap + ] | UArray UInt | UArray UReal -> - [Assignment ((decl_id, flat_type, []), readfnapp decl_var) |> swrap] - | UFun _ | UMathLibraryFunction -> - raise_s [%message "Cannot read a function type."] + [ Assignment ((decl_id, flat_type, []), readfnapp decl_var) |> swrap ] + | UFun _ | UMathLibraryFunction -> raise_s [%message "Cannot read a function type."] | UVector | URowVector | UMatrix | UArray _ -> - let decl, assign, flat_var = - let decl_id = decl_id ^ "_flat__" in - ( Stmt.Fixed.Pattern.Decl - {decl_adtype= AutoDiffable; decl_id; decl_type= Unsized flat_type} - |> swrap - , Assignment ((decl_id, flat_type, []), readfnapp decl_var) |> swrap - , { Expr.Fixed.pattern= Var decl_id - ; meta= - Expr.Typed.Meta.{loc= smeta; type_= flat_type; adlevel= DataOnly} - } ) - in - let bodyfn var = - let pos_increment = - [ Assignment ((pos, UInt, []), Expr.Helpers.(binop pos_var Plus one)) - |> swrap ] - in - let read_indexed _ = - { Expr.Fixed.pattern= Indexed (flat_var, [Single pos_var]) - ; meta= Expr.Typed.Meta.{flat_var.meta with type_= scalar} } - in - SList - ( Stmt.Helpers.assign_indexed (SizedType.to_unsized st) decl_id smeta - read_indexed var - :: pos_increment ) + let decl, assign, flat_var = + let decl_id = decl_id ^ "_flat__" in + ( Stmt.Fixed.Pattern.Decl + { decl_adtype = AutoDiffable; decl_id; decl_type = Unsized flat_type } |> swrap + , Assignment ((decl_id, flat_type, []), readfnapp decl_var) |> swrap + , { Expr.Fixed.pattern = Var decl_id + ; meta = Expr.Typed.Meta.{ loc = smeta; type_ = flat_type; adlevel = DataOnly } + } ) + in + let bodyfn var = + let pos_increment = + [ Assignment ((pos, UInt, []), Expr.Helpers.(binop pos_var Plus one)) |> swrap ] in - let pos_reset = - Stmt.Fixed.Pattern.Assignment - ((pos, UInt, []), Expr.Helpers.loop_bottom) - |> swrap + let read_indexed _ = + { Expr.Fixed.pattern = Indexed (flat_var, [ Single pos_var ]) + ; meta = Expr.Typed.Meta.{ flat_var.meta with type_ = scalar } + } in - [ Block - [ decl; assign; pos_reset - ; Stmt.Helpers.for_scalar_inv st bodyfn decl_var smeta ] - |> swrap ] + SList + (Stmt.Helpers.assign_indexed + (SizedType.to_unsized st) + decl_id + smeta + read_indexed + var + :: pos_increment) + |> swrap + in + let pos_reset = + Stmt.Fixed.Pattern.Assignment ((pos, UInt, []), Expr.Helpers.loop_bottom) |> swrap + in + [ Block + [ decl; assign; pos_reset; Stmt.Helpers.for_scalar_inv st bodyfn decl_var smeta ] + |> swrap + ] +;; let rec base_ut_to_string = function | UnsizedType.UMatrix -> "matrix" @@ -145,90 +148,93 @@ let rec base_ut_to_string = function | 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)] + | t -> raise_s [%message "Another place where it's weird to get " (t : UnsizedType.t)] +;; -let param_read smeta +let param_read + smeta ( decl_id - , Program.({ out_constrained_st= cst - ; out_unconstrained_st= ucst - ; out_block; _ }) ) = - if not (out_block = Parameters) then [] - else + , Program.{ out_constrained_st = cst; out_unconstrained_st = ucst; out_block; _ } ) + = + if not (out_block = Parameters) + then [] + else ( let decl_id, decl = match cst = ucst with - | true -> (decl_id, []) + | 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 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 } ] in let unconstrained_decl_var = let meta = - Expr.Typed.Meta.create ~loc:smeta + Expr.Typed.Meta.create + ~loc:smeta ~type_:SizedType.(to_unsized cst) - ~adlevel:AutoDiffable () + ~adlevel:AutoDiffable + () in - Expr.Fixed.{meta; pattern= Var decl_id} + 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}) + 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 + Stmt.Helpers.assign_indexed (SizedType.to_unsized cst) decl_id smeta readfnapp var in - decl @ [Stmt.Helpers.for_eigen ucst bodyfn unconstrained_decl_var smeta] + decl @ [ Stmt.Helpers.for_eigen ucst bodyfn unconstrained_decl_var smeta ]) +;; let escape_name str = str |> String.substr_replace_all ~pattern:"." ~with_:"_" |> String.substr_replace_all ~pattern:"-" ~with_:"_" +;; -let rec add_jacobians Stmt.Fixed.({meta= smeta; pattern}) = +let rec add_jacobians Stmt.Fixed.{ meta = smeta; pattern } = match pattern with - | Assignment (lhs, {pattern= FunApp (CompilerInternal, f, args); meta= emeta}) + | Assignment (lhs, { pattern = FunApp (CompilerInternal, f, args); meta = emeta }) when Internal_fun.of_string_opt f = Some FnConstrain -> - let var n = Expr.{Fixed.pattern= Var n; meta= Typed.Meta.empty} in - let assign rhs = - Stmt.{Fixed.pattern= Assignment (lhs, rhs); meta= smeta} - in - { Stmt.Fixed.pattern= - IfElse - ( var "jacobian__" - , assign - { Expr.Fixed.pattern= - FunApp (CompilerInternal, f, args @ [var "lp__"]) - ; meta= emeta } - , Some - (assign - { Expr.Fixed.pattern= FunApp (CompilerInternal, f, args) - ; meta= emeta }) ) - ; meta= smeta } - | ptn -> - Stmt.Fixed.{pattern= Pattern.map Fn.id add_jacobians ptn; meta= smeta} + let var n = Expr.{ Fixed.pattern = Var n; meta = Typed.Meta.empty } in + let assign rhs = Stmt.{ Fixed.pattern = Assignment (lhs, rhs); meta = smeta } in + { Stmt.Fixed.pattern = + IfElse + ( var "jacobian__" + , assign + { Expr.Fixed.pattern = FunApp (CompilerInternal, f, args @ [ var "lp__" ]) + ; meta = emeta + } + , Some + (assign + { Expr.Fixed.pattern = FunApp (CompilerInternal, f, args); meta = emeta }) + ) + ; meta = smeta + } + | ptn -> Stmt.Fixed.{ pattern = Pattern.map Fn.id add_jacobians ptn; meta = smeta } +;; (* Make sure that all if-while-and-for bodies are safely wrapped in a block in such a way that we can insert a location update before. The blocks make sure that the program with the inserted location update is still well-formed C++ though. *) -let rec ensure_body_in_block (Stmt.Fixed.({pattern; _}) as stmt) = +let rec ensure_body_in_block (Stmt.Fixed.{ pattern; _ } as stmt) = let in_block stmt = let pattern = Stmt.Fixed.( match stmt.pattern with | Block l | SList l -> Pattern.Block l - | _ -> Block [stmt]) + | _ -> Block [ stmt ]) in - {stmt with pattern} + { stmt with pattern } in let ensure_body_in_block_base pattern = Stmt.Fixed.Pattern.( @@ -237,51 +243,55 @@ let rec ensure_body_in_block (Stmt.Fixed.({pattern; _}) as stmt) = | _ -> pattern) in let pattern = - ensure_body_in_block_base - Stmt.Fixed.(Pattern.map Fn.id ensure_body_in_block pattern) + ensure_body_in_block_base Stmt.Fixed.(Pattern.map Fn.id ensure_body_in_block pattern) in - {stmt with pattern} + { stmt with pattern } +;; let rec flatten_slists_list ls = let flatten_slist stmt = - Stmt.Fixed.(match stmt.pattern with SList ls -> ls | _ -> [stmt]) + Stmt.Fixed.( + match stmt.pattern with + | SList ls -> ls + | _ -> [ stmt ]) in let rec flatten_slists_stmt stmt = let pattern = Stmt.Fixed.( match stmt.pattern with | Block ls -> - Pattern.Block - (List.concat_map - ~f:(Fn.compose flatten_slist flatten_slists_stmt) - ls) + Pattern.Block + (List.concat_map ~f:(Fn.compose flatten_slist flatten_slists_stmt) ls) | pattern -> Pattern.map Fn.id flatten_slists_stmt pattern) in - {stmt with pattern} + { stmt with pattern } in List.concat_map ls ~f:(fun stmt -> Stmt.Fixed.( match stmt.pattern with | SList ls -> flatten_slists_list ls - | _ -> [stmt]) ) + | _ -> [ stmt ])) |> List.map ~f:flatten_slists_stmt +;; let%expect_test "Flatten slists" = - let e pattern = Expr.Fixed.{meta= (); pattern} in - let s pattern = Stmt.Fixed.{meta= (); pattern} in + let e pattern = Expr.Fixed.{ meta = (); pattern } in + let s pattern = Stmt.Fixed.{ meta = (); pattern } in let stmt = Stmt.Fixed.Pattern.( [ SList [ Block [ SList - [ While (e (Var "hi"), Block [SList [Break |> s] |> s] |> s) - |> s ] - |> s ] - |> s ] - |> s ] + [ While (e (Var "hi"), Block [ SList [ Break |> s ] |> s ] |> s) |> s ] + |> s + ] + |> s + ] + |> s + ] |> flatten_slists_list) in - print_s [%sexp (stmt : (unit, unit) Stmt.Fixed.t list)] ; + print_s [%sexp (stmt : (unit, unit) Stmt.Fixed.t list)]; [%expect {| (((pattern @@ -291,55 +301,57 @@ let%expect_test "Flatten slists" = ((pattern (Block (((pattern Break) (meta ()))))) (meta ())))) (meta ()))))) (meta ()))) |}] +;; let add_reads vars mkread stmts = let var_names = String.Map.of_alist_exn vars in - let add_read_to_decl (Stmt.Fixed.({pattern; meta}) as stmt) = + let add_read_to_decl (Stmt.Fixed.{ pattern; meta } as stmt) = match pattern with - | Decl {decl_id; _} when Map.mem var_names decl_id -> - stmt :: mkread meta (decl_id, Map.find_exn var_names decl_id) - | _ -> [stmt] + | Decl { decl_id; _ } when Map.mem var_names decl_id -> + stmt :: mkread meta (decl_id, Map.find_exn var_names decl_id) + | _ -> [ stmt ] in List.concat_map ~f:add_read_to_decl stmts +;; let gen_write (decl_id, sizedtype) = let bodyfn var = - Stmt.Helpers.internal_nrfunapp FnWriteParam [var] Location_span.empty + Stmt.Helpers.internal_nrfunapp FnWriteParam [ var ] Location_span.empty in - let meta = - {Expr.Typed.Meta.empty with type_= SizedType.to_unsized sizedtype} - in - let expr = Expr.Fixed.{meta; pattern= Var decl_id} in + let meta = { Expr.Typed.Meta.empty with type_ = SizedType.to_unsized sizedtype } in + let expr = Expr.Fixed.{ meta; pattern = Var decl_id } in Stmt.Helpers.for_scalar_inv sizedtype bodyfn expr Location_span.empty +;; let gen_write_unconstrained (decl_id, sizedtype) = let bodyfn var = let var = match var.Expr.Fixed.pattern with - | Indexed ({pattern= Indexed (expr, idcs1); _}, idcs2) -> - {var with pattern= Indexed (expr, idcs1 @ idcs2)} + | Indexed ({ pattern = Indexed (expr, idcs1); _ }, idcs2) -> + { var with pattern = Indexed (expr, idcs1 @ idcs2) } | _ -> var in - Stmt.Helpers.internal_nrfunapp FnWriteParam [var] Location_span.empty - in - let meta = - {Expr.Typed.Meta.empty with type_= SizedType.to_unsized sizedtype} + Stmt.Helpers.internal_nrfunapp FnWriteParam [ var ] Location_span.empty in - let expr = Expr.Fixed.{meta; pattern= Var decl_id} in + let meta = { Expr.Typed.Meta.empty with type_ = SizedType.to_unsized sizedtype } in + let expr = Expr.Fixed.{ meta; pattern = Var decl_id } in let writefn var = Stmt.Helpers.for_scalar_inv (SizedType.inner_type sizedtype) - bodyfn var Location_span.empty + bodyfn + var + Location_span.empty in Stmt.Helpers.for_eigen sizedtype writefn expr Location_span.empty +;; -let rec contains_var_expr is_vident accum Expr.Fixed.({pattern; _}) = +let rec contains_var_expr is_vident accum Expr.Fixed.{ pattern; _ } = accum || match pattern with | Var v when is_vident v -> true - | pattern -> - Expr.Fixed.Pattern.fold (contains_var_expr is_vident) false pattern + | pattern -> Expr.Fixed.Pattern.fold (contains_var_expr is_vident) false pattern +;; (* When a parameter's unconstrained type and its constrained type are different, we generate a new variable "_in__" and read into that. We now need @@ -349,44 +361,35 @@ let rec contains_var_expr is_vident accum Expr.Fixed.({pattern; _}) = let constrain_in_params outvars stmts = let is_target_var = function | ( name - , { Program.out_unconstrained_st - ; out_constrained_st - ; out_block= Parameters; _ } ) - when not (out_unconstrained_st = out_constrained_st) -> - Some name + , { Program.out_unconstrained_st; out_constrained_st; out_block = Parameters; _ } ) + when not (out_unconstrained_st = out_constrained_st) -> Some name | _ -> None in - let target_vars = - List.filter_map outvars ~f:is_target_var |> String.Set.of_list - in - let rec change_constrain_target (Stmt.Fixed.({pattern; _}) as s) = + let target_vars = List.filter_map outvars ~f:is_target_var |> String.Set.of_list in + let rec change_constrain_target (Stmt.Fixed.{ pattern; _ } as s) = match pattern with - | Assignment (_, {pattern= FunApp (CompilerInternal, f, args); _}) - when ( Internal_fun.of_string_opt f = Some FnConstrain - || Internal_fun.of_string_opt f = Some FnUnconstrain ) - && List.exists args - ~f:(contains_var_expr (Set.mem target_vars) false) -> - let rec change_var_expr (Expr.Fixed.({pattern; _}) as e) = - match pattern with - | Var vident when Set.mem target_vars vident -> - {e with pattern= Var (vident ^ "_in__")} - | pattern -> - {e with pattern= Expr.Fixed.Pattern.map change_var_expr pattern} - in - let rec change_var_stmt s = - Stmt.Fixed. - { s with - pattern= Pattern.map change_var_expr change_var_stmt s.pattern } - in - change_var_stmt s - | pattern -> + | Assignment (_, { pattern = FunApp (CompilerInternal, f, args); _ }) + when (Internal_fun.of_string_opt f = Some FnConstrain + || Internal_fun.of_string_opt f = Some FnUnconstrain) + && List.exists args ~f:(contains_var_expr (Set.mem target_vars) false) -> + let rec change_var_expr (Expr.Fixed.{ pattern; _ } as e) = + match pattern with + | Var vident when Set.mem target_vars vident -> + { e with pattern = Var (vident ^ "_in__") } + | pattern -> { e with pattern = Expr.Fixed.Pattern.map change_var_expr pattern } + in + let rec change_var_stmt s = Stmt.Fixed. - {s with pattern= Pattern.map Fn.id change_constrain_target pattern} + { s with pattern = Pattern.map change_var_expr change_var_stmt s.pattern } + in + change_var_stmt s + | pattern -> + Stmt.Fixed.{ s with pattern = Pattern.map Fn.id change_constrain_target pattern } in List.map ~f:change_constrain_target stmts +;; -let fn_name_map = - String.Map.of_alist_exn [("integrate_ode", "integrate_ode_rk45")] +let fn_name_map = String.Map.of_alist_exn [ "integrate_ode", "integrate_ode_rk45" ] let rec map_fn_names s = let rec map_fn_names_expr e = @@ -394,241 +397,247 @@ let rec map_fn_names s = Expr.Fixed.( match e.pattern with | FunApp (k, f, a) when Map.mem fn_name_map f -> - Pattern.FunApp (k, Map.find_exn fn_name_map f, a) + Pattern.FunApp (k, Map.find_exn fn_name_map f, a) | expr -> Pattern.map map_fn_names_expr expr) in - {e with pattern} + { e with pattern } in let stmt = Stmt.Fixed.( match s.pattern with | NRFunApp (k, f, a) when Map.mem fn_name_map f -> - Pattern.NRFunApp (k, Map.find_exn fn_name_map f, a) + Pattern.NRFunApp (k, Map.find_exn fn_name_map f, a) | stmt -> Pattern.map map_fn_names_expr map_fn_names stmt) in - {s with pattern= stmt} + { s with pattern = stmt } +;; let rec insert_before f to_insert = function | [] -> to_insert | hd :: tl -> - if f hd then to_insert @ (hd :: tl) - else hd :: insert_before f to_insert tl + if f hd then to_insert @ (hd :: tl) else hd :: insert_before f to_insert tl +;; let is_opencl_var = String.is_suffix ~suffix:opencl_suffix -let rec collect_vars_expr is_target accum Expr.Fixed.({pattern; _}) = - Set.union accum - ( match pattern with - | Var s when is_target s -> String.Set.of_list [s] - | x -> - Expr.Fixed.Pattern.fold - (collect_vars_expr is_target) - String.Set.empty x ) +let rec collect_vars_expr is_target accum Expr.Fixed.{ pattern; _ } = + Set.union + accum + (match pattern with + | Var s when is_target s -> String.Set.of_list [ s ] + | x -> Expr.Fixed.Pattern.fold (collect_vars_expr is_target) String.Set.empty x) +;; let collect_opencl_vars s = let rec go accum s = - Stmt.Fixed.( - Pattern.fold (collect_vars_expr is_opencl_var) go accum s.pattern) + Stmt.Fixed.(Pattern.fold (collect_vars_expr is_opencl_var) go accum s.pattern) in go String.Set.empty s +;; let%expect_test "collect vars expr" = - let mkvar s = Expr.{Fixed.pattern= Var s; meta= Typed.Meta.empty} in - let args = List.map ~f:mkvar ["y"; "x_opencl__"; "z"; "w_opencl__"] in + let mkvar s = Expr.{ Fixed.pattern = Var s; meta = Typed.Meta.empty } in + let args = List.map ~f:mkvar [ "y"; "x_opencl__"; "z"; "w_opencl__" ] in let fnapp = - Expr. - {Fixed.pattern= FunApp (StanLib, "print", args); meta= Typed.Meta.empty} + Expr.{ Fixed.pattern = FunApp (StanLib, "print", args); meta = Typed.Meta.empty } in - Stmt.Fixed.{pattern= TargetPE fnapp; meta= Location_span.empty} - |> collect_opencl_vars |> String.Set.sexp_of_t |> print_s ; + Stmt.Fixed.{ pattern = TargetPE fnapp; meta = Location_span.empty } + |> collect_opencl_vars + |> String.Set.sexp_of_t + |> print_s; [%expect {| (w_opencl__ x_opencl__) |}] +;; let%expect_test "insert before" = - let l = [1; 2; 3; 4; 5; 6] |> insert_before (( = ) 6) [999] in - [%sexp (l : int list)] |> print_s ; + let l = [ 1; 2; 3; 4; 5; 6 ] |> insert_before (( = ) 6) [ 999 ] in + [%sexp (l : int list)] |> print_s; [%expect {| (1 2 3 4 5 999 6) |}] +;; let validate_sized decl_id meta transform st = let check fn x = - Stmt.Helpers.internal_nrfunapp fn - Expr.Helpers. - [str decl_id; str (Fmt.strf "%a" Expression_gen.pp_expr x); x] + Stmt.Helpers.internal_nrfunapp + fn + Expr.Helpers.[ str decl_id; str (Fmt.strf "%a" Expression_gen.pp_expr x); x ] meta in let nrfunapp fname args = - Stmt.Fixed.{pattern= NRFunApp (CompilerInternal, fname, args); meta} + Stmt.Fixed.{ pattern = NRFunApp (CompilerInternal, fname, args); meta } in let rec dims_check = function | SizedType.SInt | SReal -> [] | SArray (st, s) -> check FnValidateSize s :: dims_check st | SVector s | SRowVector s -> - let fn = - match transform with - | Some Program.Simplex -> Internal_fun.FnValidateSizeSimplex - | Some UnitVector -> FnValidateSizeUnitVector - | _ -> FnValidateSize - in - [check fn s] + let fn = + match transform with + | Some Program.Simplex -> Internal_fun.FnValidateSizeSimplex + | Some UnitVector -> FnValidateSizeUnitVector + | _ -> FnValidateSize + in + [ check fn s ] | SMatrix (rows, cols) -> - let validate_rows = - match transform with - | Some CholeskyCov -> - nrfunapp "check_greater_or_equal" - Expr.Helpers. - [ str ("cholesky_factor_cov " ^ decl_id) - ; str "num rows (must be greater or equal to num cols)" - ; rows; cols ] - | _ -> check FnValidateSize rows - in - [validate_rows; check FnValidateSize cols] + let validate_rows = + match transform with + | Some CholeskyCov -> + nrfunapp + "check_greater_or_equal" + Expr.Helpers. + [ str ("cholesky_factor_cov " ^ decl_id) + ; str "num rows (must be greater or equal to num cols)" + ; rows + ; cols + ] + | _ -> check FnValidateSize rows + in + [ validate_rows; check FnValidateSize cols ] in dims_check st +;; let rec add_validate_dims outvars stmts = let transforms = List.filter_map ~f:(function - | decl_id, Program.({out_block= Parameters; out_trans; _}) -> - Some (decl_id, out_trans) + | decl_id, Program.{ out_block = Parameters; out_trans; _ } -> + Some (decl_id, out_trans) | _ -> None) outvars |> String.Map.of_alist_exn in let with_size_checks = function - | Stmt.Fixed.({pattern= Decl {decl_id; decl_type= Sized st; _}; meta}) as - decl -> - let tr = Map.find transforms decl_id in - validate_sized decl_id meta tr st @ [decl] - | stmt -> [validate_dims_stmt stmt] + | Stmt.Fixed.{ pattern = Decl { decl_id; decl_type = Sized st; _ }; meta } as decl -> + let tr = Map.find transforms decl_id in + validate_sized decl_id meta tr st @ [ decl ] + | stmt -> [ validate_dims_stmt stmt ] in List.concat_map ~f:with_size_checks stmts and validate_dims_stmt stmt = let pattern = match stmt.pattern with - | Stmt.Fixed.Pattern.Block s -> - Stmt.Fixed.Pattern.Block (add_validate_dims [] s) + | Stmt.Fixed.Pattern.Block s -> Stmt.Fixed.Pattern.Block (add_validate_dims [] s) | SList s -> SList (add_validate_dims [] s) | While (a, b) -> While (a, validate_dims_stmt b) - | For f -> For {f with body= validate_dims_stmt f.body} + | For f -> For { f with body = validate_dims_stmt f.body } | IfElse (p, t, e) -> - IfElse (p, validate_dims_stmt t, Option.map ~f:validate_dims_stmt e) + IfElse (p, validate_dims_stmt t, Option.map ~f:validate_dims_stmt e) | s -> s in - {stmt with pattern} + { stmt with pattern } +;; let map_prog_stmt_lists f (p : ('a, 'b) Program.t) = { p with - Program.prepare_data= f p.prepare_data - ; log_prob= f p.log_prob - ; generate_quantities= f p.generate_quantities - ; transform_inits= f p.transform_inits } + Program.prepare_data = f p.prepare_data + ; log_prob = f p.log_prob + ; generate_quantities = f p.generate_quantities + ; transform_inits = f p.transform_inits + } +;; let trans_prog (p : Program.Typed.t) = let p = Program.map Fn.id map_fn_names p in let init_pos = [ Stmt.Fixed.Pattern.Decl - {decl_adtype= DataOnly; decl_id= pos; decl_type= Sized SInt} - ; Assignment ((pos, UInt, []), Expr.Helpers.loop_bottom) ] - |> List.map ~f:(fun pattern -> - Stmt.Fixed.{pattern; meta= Location_span.empty} ) + { decl_adtype = DataOnly; decl_id = pos; decl_type = Sized SInt } + ; Assignment ((pos, UInt, []), Expr.Helpers.loop_bottom) + ] + |> List.map ~f:(fun pattern -> Stmt.Fixed.{ pattern; meta = Location_span.empty }) in let get_pname_cst = function - | name, {Program.out_block= Parameters; out_constrained_st; _} -> - Some (name, out_constrained_st) + | name, { Program.out_block = Parameters; out_constrained_st; _ } -> + Some (name, out_constrained_st) | _ -> None in let get_pname_ust = function | ( name - , { Program.out_block= Parameters - ; out_constrained_st - ; out_unconstrained_st; _ } ) + , { Program.out_block = Parameters; out_constrained_st; out_unconstrained_st; _ } ) when SizedType.to_unsized out_constrained_st - = SizedType.to_unsized out_unconstrained_st -> - Some (name, out_unconstrained_st) - | name, {Program.out_block= Parameters; out_unconstrained_st; _} -> - Some (name ^ "_free__", out_unconstrained_st) + = SizedType.to_unsized out_unconstrained_st -> Some (name, out_unconstrained_st) + | name, { Program.out_block = Parameters; out_unconstrained_st; _ } -> + Some (name ^ "_free__", out_unconstrained_st) | _ -> None in let constrained_params = List.filter_map ~f:get_pname_cst p.output_vars in let free_params = List.filter_map ~f:get_pname_ust p.output_vars in let param_writes, tparam_writes, gq_writes = - List.map p.output_vars - ~f:(fun (name, {out_constrained_st= st; out_block; _}) -> - (out_block, gen_write (name, st)) ) + List.map p.output_vars ~f:(fun (name, { out_constrained_st = st; out_block; _ }) -> + out_block, gen_write (name, st)) |> List.partition3_map ~f:(fun (b, x) -> match b with | Parameters -> `Fst x | TransformedParameters -> `Snd x - | GeneratedQuantities -> `Trd x ) + | GeneratedQuantities -> `Trd x) in let tparam_start stmt = Stmt.Fixed.( match stmt.pattern with | IfElse (cond, _, _) - when contains_var_expr - (( = ) "emit_transformed_parameters__") - false cond -> - true + when contains_var_expr (( = ) "emit_transformed_parameters__") false cond -> true | _ -> false) in - let gq_start Stmt.Fixed.({pattern; _}) = + let gq_start Stmt.Fixed.{ pattern; _ } = match pattern with | IfElse - ( { pattern= - FunApp (_, _, [{pattern= Var "emit_generated_quantities__"; _}]); _ + ( { pattern = FunApp (_, _, [ { pattern = Var "emit_generated_quantities__"; _ } ]) + ; _ } , _ - , _ ) -> - true + , _ ) -> true | _ -> false in let translate_to_open_cl stmts = - if !use_opencl then - let decl Stmt.Fixed.({pattern; _}) = - match pattern with Decl d -> Some d.decl_id | _ -> None + if !use_opencl + then ( + let decl Stmt.Fixed.{ pattern; _ } = + match pattern with + | Decl d -> Some d.decl_id + | _ -> None in let data_var_idents = List.filter_map ~f:decl p.prepare_data in let switch_expr = switch_expr_to_opencl data_var_idents in let rec trans_stmt_to_opencl s = Stmt.Fixed. - { s with - pattern= Pattern.map switch_expr trans_stmt_to_opencl s.pattern } + { s with pattern = Pattern.map switch_expr trans_stmt_to_opencl s.pattern } in - List.map stmts ~f:trans_stmt_to_opencl + List.map stmts ~f:trans_stmt_to_opencl) else stmts in let functions_block = List.map - ~f:(fun def -> {def with fdbody= validate_dims_stmt def.fdbody}) + ~f:(fun def -> { def with fdbody = validate_dims_stmt def.fdbody }) p.functions_block in let tparam_writes_cond = match tparam_writes with | [] -> [] | _ -> - [ Stmt.Fixed. - { pattern= - IfElse - ( Expr. - { Fixed.pattern= Var "emit_transformed_parameters__" - ; meta= Typed.Meta.empty } - , {pattern= SList tparam_writes; meta= Location_span.empty} - , None ) - ; meta= Location_span.empty } ] + [ Stmt.Fixed. + { pattern = + IfElse + ( Expr. + { Fixed.pattern = Var "emit_transformed_parameters__" + ; meta = Typed.Meta.empty + } + , { pattern = SList tparam_writes; meta = Location_span.empty } + , None ) + ; meta = Location_span.empty + } + ] in let generate_quantities = - ( p.generate_quantities + (p.generate_quantities |> add_validate_dims p.output_vars |> add_reads p.output_vars param_read |> translate_to_open_cl |> constrain_in_params p.output_vars |> insert_before tparam_start param_writes - |> insert_before gq_start tparam_writes_cond ) + |> insert_before gq_start tparam_writes_cond) @ gq_writes in let log_prob = - p.log_prob |> List.map ~f:add_jacobians + p.log_prob + |> List.map ~f:add_jacobians |> add_validate_dims p.output_vars |> add_reads p.output_vars param_read |> constrain_in_params p.output_vars @@ -638,54 +647,51 @@ let trans_prog (p : Program.Typed.t) = String.Set.union_list (List.concat_map ~f:(List.map ~f:collect_opencl_vars) - [log_prob; generate_quantities]) + [ log_prob; generate_quantities ]) |> String.Set.to_list in let to_matrix_cl_stmts = List.concat_map opencl_vars ~f:(fun vident -> - let vident_sans_opencl = - String.chop_suffix_exn ~suffix:opencl_suffix vident - in + let vident_sans_opencl = String.chop_suffix_exn ~suffix:opencl_suffix vident in let type_of_input_var = - match - List.Assoc.find p.input_vars vident_sans_opencl ~equal:String.equal - with + match List.Assoc.find p.input_vars vident_sans_opencl ~equal:String.equal with | Some st -> SizedType.to_unsized st | None -> UnsizedType.UMatrix in [ Stmt.Fixed. - { pattern= + { pattern = Decl - { decl_adtype= DataOnly - ; decl_id= vident - ; decl_type= Type.Unsized type_of_input_var } - ; meta= Location_span.empty } - ; { pattern= + { decl_adtype = DataOnly + ; decl_id = vident + ; decl_type = Type.Unsized type_of_input_var + } + ; meta = Location_span.empty + } + ; { pattern = Assignment ( (vident, type_of_input_var, []) , to_matrix_cl - { pattern= Var vident_sans_opencl - ; meta= Expr.Typed.Meta.empty } ) - ; meta= Location_span.empty } ] ) + { pattern = Var vident_sans_opencl; meta = Expr.Typed.Meta.empty } ) + ; meta = Location_span.empty + } + ]) in let p = { p with functions_block ; log_prob - ; prog_name= escape_name p.prog_name - ; prepare_data= + ; prog_name = escape_name p.prog_name + ; prepare_data = init_pos - @ ( add_validate_dims [] p.prepare_data - |> add_reads p.input_vars data_read ) + @ (add_validate_dims [] p.prepare_data |> add_reads p.input_vars data_read) @ to_matrix_cl_stmts - ; transform_inits= + ; transform_inits = init_pos - @ ( add_validate_dims p.output_vars p.transform_inits - |> add_reads constrained_params data_read ) + @ (add_validate_dims p.output_vars p.transform_inits + |> add_reads constrained_params data_read) @ List.map ~f:gen_write_unconstrained free_params - ; generate_quantities } + ; generate_quantities + } in - Program.( - p - |> map Fn.id ensure_body_in_block - |> map_prog_stmt_lists flatten_slists_list) + Program.(p |> map Fn.id ensure_body_in_block |> map_prog_stmt_lists flatten_slists_list) +;; diff --git a/src/stan_math_backend/Transform_Mir.mli b/src/stan_math_backend/Transform_Mir.mli index 2cd077d25b..6263fbb26f 100644 --- a/src/stan_math_backend/Transform_Mir.mli +++ b/src/stan_math_backend/Transform_Mir.mli @@ -4,8 +4,8 @@ val trans_prog : Program.Typed.t -> Program.Typed.t val is_opencl_var : string -> bool val use_opencl : bool ref -val validate_sized : - string +val validate_sized + : string -> 'a -> 'b Program.transformation option -> Expr.Typed.t SizedType.t diff --git a/src/stanc/stanc.ml b/src/stanc/stanc.ml index dc5b161920..d3a78acd29 100644 --- a/src/stanc/stanc.ml +++ b/src/stanc/stanc.ml @@ -42,16 +42,13 @@ let options = , " For debugging purposes: print the parser actions" ) ; ( "--debug-ast" , Arg.Set Debugging.ast_printing - , " For debugging purposes: print the undecorated AST, before semantic \ - checking" ) + , " For debugging purposes: print the undecorated AST, before semantic checking" ) ; ( "--debug-decorated-ast" , Arg.Set Debugging.typed_ast_printing - , " For debugging purposes: print the decorated AST, after semantic \ - checking" ) + , " For debugging purposes: print the decorated AST, after semantic checking" ) ; ( "--debug-generate-data" , Arg.Set generate_data - , " For debugging purposes: generate a mock dataset to run the model on" - ) + , " For debugging purposes: generate a mock dataset to run the model on" ) ; ( "--debug-mir" , Arg.Set dump_mir , " For debugging purposes: print the MIR as an S-expression." ) @@ -60,24 +57,22 @@ let options = , " For debugging purposes: pretty-print the MIR." ) ; ( "--debug-optimized-mir" , Arg.Set dump_opt_mir - , " For debugging purposes: print the MIR after it's been optimized. \ - Only has an effect when optimizations are turned on." ) + , " For debugging purposes: print the MIR after it's been optimized. Only has an \ + effect when optimizations are turned on." ) ; ( "--debug-optimized-mir-pretty" , Arg.Set dump_opt_mir_pretty - , " For debugging purposes: pretty print the MIR after it's been \ - optimized. Only has an effect when optimizations are turned on." ) + , " For debugging purposes: pretty print the MIR after it's been optimized. Only \ + has an effect when optimizations are turned on." ) ; ( "--debug-transformed-mir" , Arg.Set dump_tx_mir - , " For debugging purposes: print the MIR after the backend has \ - transformed it." ) + , " For debugging purposes: print the MIR after the backend has transformed it." ) ; ( "--debug-transformed-mir-pretty" , Arg.Set dump_tx_mir_pretty - , " For debugging purposes: pretty print the MIR after the backend has \ - transformed it." ) + , " For debugging purposes: pretty print the MIR after the backend has transformed \ + it." ) ; ( "--dump-stan-math-signatures" , Arg.Set dump_stan_math_sigs - , "Dump out the list of supported type signatures for Stan Math backend." - ) + , "Dump out the list of supported type signatures for Stan Math backend." ) ; ( "--warn-uninitialized" , Arg.Set warn_uninitialized , " Emit warnings about uninitialized variables to stderr. Currently an \ @@ -94,13 +89,12 @@ let options = ; ( "--version" , Arg.Unit (fun _ -> - print_endline (version ^ " " ^ "(" ^ Sys.os_type ^ ")") ; - exit 0 ) + print_endline (version ^ " " ^ "(" ^ Sys.os_type ^ ")"); + exit 0) , " Display stanc version number" ) ; ( "--name" , Arg.Set_string Semantic_check.model_name - , " Take a string to set the model name (default = \ - \"$model_filename_model\")" ) + , " Take a string to set the model name (default = \"$model_filename_model\")" ) ; ( "--O" , Arg.Set optimize , "Allow the compiler to apply all optimizations to the Stan code." ) @@ -119,110 +113,117 @@ let options = , " Deprecated. Same as --allow-undefined." ) ; ( "--include-paths" , Arg.String - (fun str -> - Preprocessor.include_paths := String.split_on_chars ~on:[','] str - ) - , " Takes a comma-separated list of directories that may contain a file \ - in an #include directive (default = \"\")" ) + (fun str -> Preprocessor.include_paths := String.split_on_chars ~on:[ ',' ] str) + , " Takes a comma-separated list of directories that may contain a file in an \ + #include directive (default = \"\")" ) ; ( "--include_paths" , Arg.String (fun str -> - Preprocessor.include_paths := - !Preprocessor.include_paths @ String.split_on_chars ~on:[','] str - ) + Preprocessor.include_paths + := !Preprocessor.include_paths @ String.split_on_chars ~on:[ ',' ] str) , " Deprecated. Same as --include-paths." ) ; ( "--use-opencl" , Arg.Set Transform_Mir.use_opencl - , " If set, try to use matrix_cl signatures." ) ] + , " If set, try to use matrix_cl signatures." ) + ] +;; let print_deprecated_arg_warning = (* is_prefix is used to also cover the --include-paths=... *) let arg_is_used arg = Array.mem ~equal:(fun x y -> String.is_prefix ~prefix:x y) Sys.argv arg in - if arg_is_used "--allow_undefined" then - eprintf "--allow_undefined is deprecated. Please use --allow-undefined.\n" ; - if arg_is_used "--include_paths" then - eprintf "--include_paths is deprecated. Please use --include-paths.\n" + if arg_is_used "--allow_undefined" + then eprintf "--allow_undefined is deprecated. Please use --allow-undefined.\n"; + if arg_is_used "--include_paths" + then eprintf "--include_paths is deprecated. Please use --include-paths.\n" +;; let model_file_err () = - Arg.usage options ("Please specify one model_file.\n\n" ^ usage) ; + Arg.usage options ("Please specify one model_file.\n\n" ^ usage); exit 127 +;; let model_file_start_char_err () = - eprintf "%s" - "Model name must not start with a number or symbol other than underscore.\n" ; + eprintf + "%s" + "Model name must not start with a number or symbol other than underscore.\n"; exit 127 +;; let add_file filename = if !model_file = "" then model_file := filename else model_file_err () +;; (** ad directives from the given file. *) let use_file filename = let ast = - if !canonicalize_program then + if !canonicalize_program + then Canonicalize.repair_syntax (Errors.without_warnings Frontend_utils.get_ast_or_exit filename) else Frontend_utils.get_ast_or_exit filename in - Debugging.ast_logger ast ; - if !pretty_print_program then - print_endline (Pretty_printing.pretty_print_program ast) ; + Debugging.ast_logger ast; + if !pretty_print_program then print_endline (Pretty_printing.pretty_print_program ast); let typed_ast = Frontend_utils.type_ast_or_exit ast in - if !canonicalize_program then + if !canonicalize_program + then print_endline (Pretty_printing.pretty_print_typed_program - (Canonicalize.canonicalize_program typed_ast)) ; - if !generate_data then - print_endline (Debug_data_generation.print_data_prog typed_ast) ; - Debugging.typed_ast_logger typed_ast ; - if not (!pretty_print_program || !canonicalize_program) then ( + (Canonicalize.canonicalize_program typed_ast)); + if !generate_data then print_endline (Debug_data_generation.print_data_prog typed_ast); + Debugging.typed_ast_logger typed_ast; + if not (!pretty_print_program || !canonicalize_program) + then ( let mir = Ast_to_Mir.trans_prog filename typed_ast in - if !dump_mir then - Sexp.pp_hum Format.std_formatter [%sexp (mir : Middle.Program.Typed.t)] ; - if !dump_mir_pretty then Program.Typed.pp Format.std_formatter mir ; - if !warn_pedantic then Pedantic_analysis.print_warn_pedantic mir ; - if !warn_uninitialized then Pedantic_analysis.print_warn_uninitialized mir ; + if !dump_mir + then Sexp.pp_hum Format.std_formatter [%sexp (mir : Middle.Program.Typed.t)]; + if !dump_mir_pretty then Program.Typed.pp Format.std_formatter mir; + if !warn_pedantic then Pedantic_analysis.print_warn_pedantic mir; + if !warn_uninitialized then Pedantic_analysis.print_warn_uninitialized mir; let tx_mir = Transform_Mir.trans_prog mir in - if !dump_tx_mir then - Sexp.pp_hum Format.std_formatter - [%sexp (tx_mir : Middle.Program.Typed.t)] ; - if !dump_tx_mir_pretty then Program.Typed.pp Format.std_formatter tx_mir ; + if !dump_tx_mir + then Sexp.pp_hum Format.std_formatter [%sexp (tx_mir : Middle.Program.Typed.t)]; + if !dump_tx_mir_pretty then Program.Typed.pp Format.std_formatter tx_mir; let opt_mir = - if !optimize then ( + if !optimize + then ( let opt = Optimize.optimization_suite tx_mir in - if !dump_opt_mir then - Sexp.pp_hum Format.std_formatter - [%sexp (opt : Middle.Program.Typed.t)] ; - if !dump_opt_mir_pretty then Program.Typed.pp Format.std_formatter opt ; - opt ) + if !dump_opt_mir + then Sexp.pp_hum Format.std_formatter [%sexp (opt : Middle.Program.Typed.t)]; + if !dump_opt_mir_pretty then Program.Typed.pp Format.std_formatter opt; + opt) else tx_mir in let cpp = Fmt.strf "%a" Stan_math_code_gen.pp_prog opt_mir in - Out_channel.write_all !output_file ~data:cpp ; - if !print_model_cpp then print_endline cpp ) + Out_channel.write_all !output_file ~data:cpp; + if !print_model_cpp then print_endline cpp) +;; let remove_dotstan s = String.drop_suffix s 5 let model_name_check_regex = Str.regexp "^[a-zA-Z_].*$" let main () = (* Parse the arguments. *) - Arg.parse options add_file usage ; - print_deprecated_arg_warning ; + Arg.parse options add_file usage; + print_deprecated_arg_warning; (* print_deprecated_arg_warning options; *) (* Deal with multiple modalities *) - if !dump_stan_math_sigs then ( - Stan_math_signatures.pretty_print_all_math_sigs Format.std_formatter () ; - exit 0 ) ; + if !dump_stan_math_sigs + then ( + Stan_math_signatures.pretty_print_all_math_sigs Format.std_formatter (); + exit 0); (* Just translate a stan program *) - if !model_file = "" then model_file_err () ; - if !Semantic_check.model_name = "" then - Semantic_check.model_name := - remove_dotstan List.(hd_exn (rev (String.split !model_file ~on:'/'))) - ^ "_model" ; + if !model_file = "" then model_file_err (); + if !Semantic_check.model_name = "" + then + Semantic_check.model_name + := remove_dotstan List.(hd_exn (rev (String.split !model_file ~on:'/'))) ^ "_model"; if not (Str.string_match model_name_check_regex !Semantic_check.model_name 0) - then model_file_start_char_err () ; - if !output_file = "" then output_file := remove_dotstan !model_file ^ ".hpp" ; + then model_file_start_char_err (); + if !output_file = "" then output_file := remove_dotstan !model_file ^ ".hpp"; use_file !model_file +;; let () = main () diff --git a/src/stancjs/stancjs.ml b/src/stancjs/stancjs.ml index 94a1d30f7a..f9474f0909 100644 --- a/src/stancjs/stancjs.ml +++ b/src/stancjs/stancjs.ml @@ -5,63 +5,67 @@ open Analysis_and_optimization open Middle open Js_of_ocaml -let print_warn_uninitialized - (uninit_vars : (Location_span.t * string) Set.Poly.t) = +let print_warn_uninitialized (uninit_vars : (Location_span.t * string) Set.Poly.t) = let show_var_info (span, var_name) = Location_span.to_string span - ^ ":\n" ^ " Warning: The variable '" ^ var_name + ^ ":\n" + ^ " Warning: The variable '" + ^ var_name ^ "' may not have been initialized.\n" in let filtered_uninit_vars = Set.filter ~f:(fun (span, _) -> span <> Location_span.empty) uninit_vars in Set.iter filtered_uninit_vars ~f:(fun v_info -> - Out_channel.output_string stderr (show_var_info v_info) ) + Out_channel.output_string stderr (show_var_info v_info)) +;; let stan2cpp model_name model_string = - Semantic_check.model_name := model_name ; + Semantic_check.model_name := model_name; let ast = Parse.parse_string Parser.Incremental.program model_string |> Result.map_error ~f:(Fmt.to_to_string Errors.pp_syntax_error) in let semantic_err_to_string = function | Result.Error (error :: _) -> - let loc = Semantic_error.location error - and msg = (Fmt.to_to_string Semantic_error.pp) error in - Result.Error (Fmt.strf "%a" Errors.pp_semantic_error (msg, loc)) + let loc = Semantic_error.location error + and msg = (Fmt.to_to_string Semantic_error.pp) error in + Result.Error (Fmt.strf "%a" Errors.pp_semantic_error (msg, loc)) | Result.Ok _ as ok -> ok | Result.Error [] -> - Result.Error - "Semantic check failed but reported no errors. This should never \ - happen." + Result.Error + "Semantic check failed but reported no errors. This should never happen." in - Result.bind ast - ~f: - (Fn.compose semantic_err_to_string Semantic_check.semantic_check_program) + Result.bind + ast + ~f:(Fn.compose semantic_err_to_string Semantic_check.semantic_check_program) |> Result.map ~f:(fun typed_ast -> let mir = Ast_to_Mir.trans_prog model_name typed_ast in - let uninitialized_vars = - Dependence_analysis.mir_uninitialized_variables mir - in - print_warn_uninitialized uninitialized_vars ; + let uninitialized_vars = Dependence_analysis.mir_uninitialized_variables mir in + print_warn_uninitialized uninitialized_vars; let tx_mir = Transform_Mir.trans_prog mir in let cpp = Fmt.strf "%a" Stan_math_code_gen.pp_prog tx_mir in - cpp ) + cpp) +;; let wrap_result = function | Result.Ok s -> - Js.Unsafe.obj - [| ("result", Js.Unsafe.inject (Js.string s)) - ; ("warnings", Js.Unsafe.inject Js.array_empty) |] + Js.Unsafe.obj + [| "result", Js.Unsafe.inject (Js.string s) + ; "warnings", Js.Unsafe.inject Js.array_empty + |] | Error e -> - Js.Unsafe.obj - [| ("errors", Js.Unsafe.inject (Array.map ~f:Js.string [|e|])) - ; ("warnings", Js.Unsafe.inject Js.array_empty) |] + Js.Unsafe.obj + [| "errors", Js.Unsafe.inject (Array.map ~f:Js.string [| e |]) + ; "warnings", Js.Unsafe.inject Js.array_empty + |] +;; -let map2 f (x, y) = (f x, f y) +let map2 f (x, y) = f x, f y let wrap2 f s1 s2 = let s1, s2 = map2 Js.to_string (s1, s2) in f s1 s2 |> wrap_result +;; let () = Js.export "stanc" (wrap2 stan2cpp) diff --git a/src/tfp_backend/Code_gen.ml b/src/tfp_backend/Code_gen.ml index dcfc3777bd..460ae74be9 100644 --- a/src/tfp_backend/Code_gen.ml +++ b/src/tfp_backend/Code_gen.ml @@ -2,10 +2,14 @@ open Core_kernel open Middle open Fmt -let is_multi_index = function Index.MultiIndex _ -> true | _ -> false +let is_multi_index = function + | Index.MultiIndex _ -> true + | _ -> false +;; let pp_call ppf (name, pp_arg, args) = pf ppf "%s(@[%a@])" name (list ~sep:comma pp_arg) args +;; let pp_call_str ppf (name, args) = pp_call ppf (name, string, args) @@ -13,43 +17,47 @@ let pystring_of_operator = function | Operator.IntDivide -> "//" | Operator.Pow -> "**" | x -> strf "%a" Operator.pp x +;; -let rec pp_expr ppf {Expr.Fixed.pattern; _} = +let rec pp_expr ppf { Expr.Fixed.pattern; _ } = match pattern with | Var ident -> string ppf ident | Lit (Str, s) -> pf ppf "%S" s | Lit (_, s) -> pf ppf "tf__.cast(%s, tf__.float64)" s | FunApp (StanLib, f, obs :: dist_params) when f = Transform_mir.dist_prefix ^ "CholeskyLKJ" -> - pf ppf "%s(@[(%a).shape[0], %a@]).log_prob(%a)" f pp_expr obs - (list ~sep:comma pp_expr) dist_params pp_expr obs + pf + ppf + "%s(@[(%a).shape[0], %a@]).log_prob(%a)" + f + pp_expr + obs + (list ~sep:comma pp_expr) + dist_params + pp_expr + obs | FunApp (StanLib, f, obs :: dist_params) when String.is_prefix ~prefix:Transform_mir.dist_prefix f -> - pf ppf "%a.log_prob(%a)" pp_call (f, pp_expr, dist_params) pp_expr obs - | FunApp (StanLib, f, args) when Operator.of_string_opt f |> Option.is_some - -> ( - match - ( Operator.of_string_opt f |> Option.value_exn |> pystring_of_operator - , args ) - with - | op, [lhs; rhs] -> pf ppf "%a %s %a" pp_paren lhs op pp_paren rhs - | op, [unary] -> pf ppf "%s%a" op pp_paren unary - | op, args -> - raise_s [%message "Need to implement" op (args : Expr.Typed.t list)] ) + pf ppf "%a.log_prob(%a)" pp_call (f, pp_expr, dist_params) pp_expr obs + | FunApp (StanLib, f, args) when Operator.of_string_opt f |> Option.is_some -> + (match Operator.of_string_opt f |> Option.value_exn |> pystring_of_operator, args with + | op, [ lhs; rhs ] -> pf ppf "%a %s %a" pp_paren lhs op pp_paren rhs + | op, [ unary ] -> pf ppf "%s%a" op pp_paren unary + | op, args -> raise_s [%message "Need to implement" op (args : Expr.Typed.t list)]) | FunApp (_, fname, args) -> pp_call ppf (fname, pp_expr, args) | TernaryIf (cond, iftrue, iffalse) -> - pf ppf "%a if %a else %a" pp_paren iftrue pp_paren cond pp_paren iffalse + pf ppf "%a if %a else %a" pp_paren iftrue pp_paren cond pp_paren iffalse | EAnd (a, b) -> pf ppf "%a and %a" pp_paren a pp_paren b | EOr (a, b) -> pf ppf "%a or %a" pp_paren a pp_paren b | Indexed (_, indices) when List.exists ~f:is_multi_index indices -> - (* + (* TF indexing options: - * tf.slice - * tf.gather - * tf.gather_nd - * tf.strided_slice -*) - raise_s [%message "Multi-indices not supported yet"] + * tf.slice + * tf.gather + * tf.gather_nd + * tf.strided_slice + *) + raise_s [%message "Multi-indices not supported yet"] | Indexed (obj, indices) -> pf ppf "%a%a" pp_expr obj pp_indices indices and pp_indices ppf = function @@ -60,20 +68,20 @@ and pp_paren ppf expr = match expr.Expr.Fixed.pattern with | TernaryIf _ | EAnd _ | EOr _ -> pf ppf "(%a)" pp_expr expr | FunApp (StanLib, f, _) when Operator.of_string_opt f |> Option.is_some -> - pf ppf "(%a)" pp_expr expr + pf ppf "(%a)" pp_expr expr | _ -> pp_expr ppf expr +;; let rec pp_stmt ppf s = match s.Stmt.Fixed.pattern with | Assignment ((lhs, _, indices), rhs) -> - pf ppf "%s%a = %a" lhs pp_indices indices pp_expr rhs + pf ppf "%s%a = %a" lhs pp_indices indices pp_expr rhs | TargetPE rhs -> pf ppf "target += tf__.reduce_sum(%a)" pp_expr rhs | NRFunApp (StanLib, f, args) | NRFunApp (UserDefined, f, args) -> - pp_call ppf (f, pp_expr, args) + pp_call ppf (f, pp_expr, args) | Break -> pf ppf "break" | Continue -> pf ppf "continue" - | Return rhs -> - pf ppf "return %a" (option ~none:(const string "None") pp_expr) rhs + | Return rhs -> pf ppf "return %a" (option ~none:(const string "None") pp_expr) rhs | Block ls | SList ls -> (list ~sep:cut pp_stmt) ppf ls | Skip -> () (* | Decl {decl_adtype= AutoDiffable; decl_id; _} -> @@ -83,44 +91,48 @@ let rec pp_stmt ppf s = their arguments. I think these functions need to be named and defined inline in general because lambdas are limited. *) - | IfElse (_, _, _) | While (_, _) | For _ | NRFunApp (CompilerInternal, _, _) - -> - raise_s [%message "Not implemented" (s : Stmt.Located.t)] + | IfElse (_, _, _) | While (_, _) | For _ | NRFunApp (CompilerInternal, _, _) -> + raise_s [%message "Not implemented" (s : Stmt.Located.t)] +;; let pp_method ppf name params intro ?(outro = []) ppbody = - pf ppf "@[def %a:@," pp_call_str (name, params) ; - (list ~sep:cut string) ppf (intro @ [""]) ; - ppbody ppf ; - if not (List.is_empty outro) then pf ppf "@ %a" (list ~sep:cut string) outro ; + pf ppf "@[def %a:@," pp_call_str (name, params); + (list ~sep:cut string) ppf (intro @ [ "" ]); + ppbody ppf; + if not (List.is_empty outro) then pf ppf "@ %a" (list ~sep:cut string) outro; pf ppf "@, @]" +;; let rec pp_cast prefix ppf (name, st) = match st with | SizedType.SArray (t, _) -> pp_cast prefix ppf (name, t) | SInt -> pf ppf "%s%s" prefix name | _ -> pf ppf "tf__.cast(%a, tf__.float64)" (pp_cast prefix) (name, SInt) +;; let pp_init ppf p = - let pp_save_data ppf (name, st) = - pf ppf "self.%s = %a" name (pp_cast "") (name, st) - in + let pp_save_data ppf (name, st) = pf ppf "self.%s = %a" name (pp_cast "") (name, st) in let pp_prep_data_stmt ppf st = pf ppf "self.%a" pp_stmt st in let ppbody ppf = match p.Program.input_vars with | [] -> pf ppf "pass" | _ -> - pf ppf "@[%a@,%a@]" - (list ~sep:cut pp_save_data) - p.Program.input_vars - (list ~sep:cut pp_prep_data_stmt) - p.Program.prepare_data + pf + ppf + "@[%a@,%a@]" + (list ~sep:cut pp_save_data) + p.Program.input_vars + (list ~sep:cut pp_prep_data_stmt) + p.Program.prepare_data in pp_method ppf "__init__" ("self" :: List.map ~f:fst p.input_vars) [] ppbody +;; let pp_var_assignment ppf s = pf ppf "%s = self.%s" s s let pp_extract_data ppf p = (list ~sep:cut pp_var_assignment) ppf (List.map ~f:fst p.Program.input_vars) +;; let pp_extract_transf_data ppf p = let extract_arg_names x = @@ -128,138 +140,167 @@ let pp_extract_transf_data ppf p = | Assignment ((lhs, _, _), _) -> Some lhs | _ -> None in - let arg_names = - List.filter_map ~f:extract_arg_names p.Program.prepare_data - in + let arg_names = List.filter_map ~f:extract_arg_names p.Program.prepare_data in (list ~sep:cut pp_var_assignment) ppf arg_names +;; let pp_log_prob_one_chain ppf p = let pp_extract_param ppf (idx, name) = pf ppf "%s = tf__.cast(params[%d], tf__.float64)" name idx in let grab_params idx = function - | name, {Program.out_block= Parameters; _} -> [(idx, name)] + | name, { Program.out_block = Parameters; _ } -> [ idx, name ] | _ -> [] in let ppbody ppf = - pf ppf "@,%s@,%a@,@,%s@,%a@,@,%s@,%a@,@,%s@,%a" "# Data" pp_extract_data p - "# Transformed data" pp_extract_transf_data p "# Parameters" + pf + ppf + "@,%s@,%a@,@,%s@,%a@,@,%s@,%a@,@,%s@,%a" + "# Data" + pp_extract_data + p + "# Transformed data" + pp_extract_transf_data + p + "# Parameters" (list ~sep:cut pp_extract_param) List.(concat (mapi p.output_vars ~f:grab_params)) - "# Target log probability computation" (list ~sep:cut pp_stmt) p.log_prob + "# Target log probability computation" + (list ~sep:cut pp_stmt) + p.log_prob in - let intro = ["target = 0"] in - let outro = ["return target"] in - pp_method ppf "log_prob_one_chain" ["self"; "params"] intro ~outro ppbody + let intro = [ "target = 0" ] in + let outro = [ "return target" ] in + pp_method ppf "log_prob_one_chain" [ "self"; "params" ] intro ~outro ppbody +;; let rec get_vident_exn e = match e.Expr.Fixed.pattern with | Var s -> s | Indexed (e, _) -> get_vident_exn e | _ -> raise_s [%message "No vident in" (e : Expr.Typed.t)] +;; -let rec contains_var_expr is_vident accum {Expr.Fixed.pattern; _} = +let rec contains_var_expr is_vident accum { Expr.Fixed.pattern; _ } = accum || match pattern with | Var v when is_vident v -> true | _ -> Expr.Fixed.Pattern.fold (contains_var_expr is_vident) false pattern +;; -let rec contains_var_stmt is_vident accum {Stmt.Fixed.pattern; _} = +let rec contains_var_stmt is_vident accum { Stmt.Fixed.pattern; _ } = Stmt.Fixed.Pattern.fold (contains_var_expr is_vident) (contains_var_stmt is_vident) - accum pattern + accum + pattern +;; let get_param_st p var = - let {Program.out_constrained_st= st; _} = + let { Program.out_constrained_st = st; _ } = List.Assoc.find_exn ~equal:( = ) p.Program.output_vars (get_vident_exn var) in st +;; let pp_log_prob ppf p = - pf ppf "@ %a@ " pp_log_prob_one_chain p ; - let intro = - ["return tf__.vectorized_map(self.log_prob_one_chain, params)"] - in - pp_method ppf "log_prob" ["self"; "params"] intro (fun _ -> ()) + pf ppf "@ %a@ " pp_log_prob_one_chain p; + let intro = [ "return tf__.vectorized_map(self.log_prob_one_chain, params)" ] in + pp_method ppf "log_prob" [ "self"; "params" ] intro (fun _ -> ()) +;; let get_params p = List.filter - ~f:(function _, {Program.out_block= Parameters; _} -> true | _ -> false) + ~f:(function + | _, { Program.out_block = Parameters; _ } -> true + | _ -> false) p.Program.output_vars +;; let pp_shapes ppf p = - let pp_shape ppf (_, {Program.out_unconstrained_st; _}) = + let pp_shape ppf (_, { Program.out_unconstrained_st; _ }) = let cast_expr ppf e = pf ppf "tf__.cast(%a, tf__.int32)" pp_expr e in - pf ppf "(nchains__, @[%a@])" + pf + ppf + "(nchains__, @[%a@])" (list ~sep:comma cast_expr) (SizedType.get_dims out_unconstrained_st) in let ppbody ppf = - pf ppf "%a@ " pp_extract_data p ; + pf ppf "%a@ " pp_extract_data p; pf ppf "return [@[%a@]]" (list ~sep:comma pp_shape) (get_params p) in - pp_method ppf "parameter_shapes" ["self"; "nchains__"] [] ppbody + pp_method ppf "parameter_shapes" [ "self"; "nchains__" ] [] ppbody +;; let pp_bijector ppf trans = let pp_call_expr ppf (name, args) = pp_call ppf (name, pp_expr, args) in let components = match trans with | Program.Identity -> [] - | Lower lb -> [("Exp", []); ("Shift", [lb])] - | Upper ub -> - [("Exp", []); ("Scale", [Expr.Helpers.float (-1.)]); ("Shift", [ub])] - | LowerUpper (lb, ub) -> [("Sigmoid", [lb; ub])] - | Offset o -> [("Shift", [o])] - | Multiplier m -> [("Scale", [m])] - | OffsetMultiplier (o, m) -> [("Scale", [m]); ("Shift", [o])] - | CholeskyCorr -> [("CorrelationCholesky", [])] - | Correlation -> [("CorrelationCholesky", []); ("CholeskyOuterProduct", [])] - | _ -> - raise_s - [%message - "Unsupported " (trans : Expr.Typed.t Program.transformation)] + | Lower lb -> [ "Exp", []; "Shift", [ lb ] ] + | Upper ub -> [ "Exp", []; "Scale", [ Expr.Helpers.float (-1.) ]; "Shift", [ ub ] ] + | LowerUpper (lb, ub) -> [ "Sigmoid", [ lb; ub ] ] + | Offset o -> [ "Shift", [ o ] ] + | Multiplier m -> [ "Scale", [ m ] ] + | OffsetMultiplier (o, m) -> [ "Scale", [ m ]; "Shift", [ o ] ] + | CholeskyCorr -> [ "CorrelationCholesky", [] ] + | Correlation -> [ "CorrelationCholesky", []; "CholeskyOuterProduct", [] ] + | _ -> raise_s [%message "Unsupported " (trans : Expr.Typed.t Program.transformation)] in match components with | [] -> pf ppf "tfb__.Identity()" | ls -> - pf ppf "tfb__.Chain([@[%a@]])" - (list ~sep:comma pp_call_expr) - List.(rev (map ls ~f:(fun (s, args) -> ("tfb__." ^ s, args)))) + pf + ppf + "tfb__.Chain([@[%a@]])" + (list ~sep:comma pp_call_expr) + List.(rev (map ls ~f:(fun (s, args) -> "tfb__." ^ s, args))) +;; let pp_bijectors ppf p = let ppbody ppf = - pf ppf "%a@ " pp_extract_data p ; - pf ppf "return [@[%a@]]" + pf ppf "%a@ " pp_extract_data p; + pf + ppf + "return [@[%a@]]" (list ~sep:comma pp_bijector) - (List.map ~f:(fun (_, {out_trans; _}) -> out_trans) (get_params p)) + (List.map ~f:(fun (_, { out_trans; _ }) -> out_trans) (get_params p)) in - pp_method ppf "parameter_bijectors" ["self"] [] ppbody + pp_method ppf "parameter_bijectors" [ "self" ] [] ppbody +;; let pp_param_names ppf p = let param_names = List.filter_map - ~f:(function name, {out_block= Parameters; _} -> Some name | _ -> None) + ~f:(function + | name, { out_block = Parameters; _ } -> Some name + | _ -> None) p.Program.output_vars in let ppbody ppf = pf ppf "return [@[%a@]]" (list ~sep:comma (fmt "%S")) param_names in - pp_method ppf "parameter_names" ["self"] [] ppbody + pp_method ppf "parameter_names" [ "self" ] [] ppbody +;; let pp_methods ppf p = - pf ppf "@ %a" pp_init p ; - pf ppf "@ %a" pp_log_prob p ; - pf ppf "@ %a" pp_shapes p ; - pf ppf "@ %a" pp_bijectors p ; + pf ppf "@ %a" pp_init p; + pf ppf "@ %a" pp_log_prob p; + pf ppf "@ %a" pp_shapes p; + pf ppf "@ %a" pp_bijectors p; pf ppf "@ %a" pp_param_names p +;; -let pp_fundef ppf {Program.fdname; fdargs; fdbody; _} = - pp_method ppf fdname +let pp_fundef ppf { Program.fdname; fdargs; fdbody; _ } = + pp_method + ppf + fdname (List.map ~f:(fun (_, name, _) -> name) fdargs) [] (fun ppf -> pp_stmt ppf fdbody) +;; let imports = {| @@ -270,11 +311,20 @@ tfd__ = tfp__.distributions tfb__ = tfp__.bijectors from tensorflow.python.ops.parallel_for import pfor as pfor__ |} +;; let pp_prog ppf (p : Program.Typed.t) = - pf ppf "@[%s@,%a@,class %s(tfd__.Distribution):@,@[%a@]@]" imports - (list ~sep:cut pp_fundef) p.functions_block p.prog_name pp_methods p ; + pf + ppf + "@[%s@,%a@,class %s(tfd__.Distribution):@,@[%a@]@]" + imports + (list ~sep:cut pp_fundef) + p.functions_block + p.prog_name + pp_methods + p; pf ppf "@ model = %s" p.prog_name +;; (* Major work to do: 1. Work awareness of distributions and bijectors into the type system diff --git a/src/tfp_backend/Transform_mir.ml b/src/tfp_backend/Transform_mir.ml index 7462362bf7..35df4a9838 100644 --- a/src/tfp_backend/Transform_mir.ml +++ b/src/tfp_backend/Transform_mir.ml @@ -6,91 +6,124 @@ let kwrds_suffix = "__" let append_kwrds_suffix x = let x_with_suffix = x ^ kwrds_suffix in - Fmt.epr "Identifier %s is a reserved word in python, renamed to %s@," x - x_with_suffix ; + Fmt.epr "Identifier %s is a reserved word in python, renamed to %s@," x x_with_suffix; x_with_suffix +;; let python_kwrds = String.Set.of_list - [ "False"; "None"; "True"; "and"; "as"; "assert"; "break"; "class" - ; "continue"; "def"; "del"; "elif"; "else"; "except"; "finally"; "for" - ; "from"; "global"; "if"; "import"; "in"; "is"; "lambda"; "nonlocal"; "not" - ; "or"; "pass"; "raise"; "return"; "try"; "while"; "with"; "yield"; "await" - ; "async" ] + [ "False" + ; "None" + ; "True" + ; "and" + ; "as" + ; "assert" + ; "break" + ; "class" + ; "continue" + ; "def" + ; "del" + ; "elif" + ; "else" + ; "except" + ; "finally" + ; "for" + ; "from" + ; "global" + ; "if" + ; "import" + ; "in" + ; "is" + ; "lambda" + ; "nonlocal" + ; "not" + ; "or" + ; "pass" + ; "raise" + ; "return" + ; "try" + ; "while" + ; "with" + ; "yield" + ; "await" + ; "async" + ] +;; -let add_suffix_to_kwrds s = - if Set.mem python_kwrds s then append_kwrds_suffix s else s +let add_suffix_to_kwrds s = if Set.mem python_kwrds s then append_kwrds_suffix s else s let remove_stan_dist_suffix s = let s = Utils.stdlib_distribution_name s in List.filter_map - (("_rng" :: Utils.distribution_suffices) @ [""]) + (("_rng" :: Utils.distribution_suffices) @ [ "" ]) ~f:(fun suffix -> String.chop_suffix ~suffix s) |> List.hd_exn +;; let capitalize_fnames = String.Set.of_list - ["normal"; "cauchy"; "gumbel"; "exponential"; "gamma"; "beta"; "poisson"] + [ "normal"; "cauchy"; "gumbel"; "exponential"; "gamma"; "beta"; "poisson" ] +;; let map_functions fname args = let open Expr in - let none = {Fixed.pattern= Var "None"; meta= Typed.Meta.empty} in - match (fname, args) with - | "multi_normal_cholesky", _ -> ("MultivariateNormalTriL", args) - | "student_t", _ -> ("StudentT", args) - | "double_exponential", _ -> ("Laplace", args) - | "lognormal", _ -> ("LogNormal", args) - | "chi_square", _ -> ("Chi2", args) - | "inv_gamma", _ -> ("InverseGamma", args) - | "lkj_corr_cholesky", _ -> ("CholeskyLKJ", args) - | "binomial_logit", _ -> ("Binomial", args) - | "bernoulli_logit", _ -> ("Bernoulli", args) - | "von_mises", _ -> ("VonMises", args) - | "binomial", [y; n; p] -> ("Binomial", [y; n; none; p]) - | "bernoulli", [y; p] -> ("Bernoulli", [y; none; p]) - | "poisson_log", [y; log_lambda] -> ("Poisson", [y; none; log_lambda]) - | "pareto", [y; y_min; alpha] -> ("Pareto", [y; alpha; y_min]) - | "neg_binomial", [y; a; b] -> - ( "NegativeBinomial" - , [y; a; none; Helpers.(binop (int 1) Divide (binop (int 1) Plus b))] ) + let none = { Fixed.pattern = Var "None"; meta = Typed.Meta.empty } in + match fname, args with + | "multi_normal_cholesky", _ -> "MultivariateNormalTriL", args + | "student_t", _ -> "StudentT", args + | "double_exponential", _ -> "Laplace", args + | "lognormal", _ -> "LogNormal", args + | "chi_square", _ -> "Chi2", args + | "inv_gamma", _ -> "InverseGamma", args + | "lkj_corr_cholesky", _ -> "CholeskyLKJ", args + | "binomial_logit", _ -> "Binomial", args + | "bernoulli_logit", _ -> "Bernoulli", args + | "von_mises", _ -> "VonMises", args + | "binomial", [ y; n; p ] -> "Binomial", [ y; n; none; p ] + | "bernoulli", [ y; p ] -> "Bernoulli", [ y; none; p ] + | "poisson_log", [ y; log_lambda ] -> "Poisson", [ y; none; log_lambda ] + | "pareto", [ y; y_min; alpha ] -> "Pareto", [ y; alpha; y_min ] + | "neg_binomial", [ y; a; b ] -> + ( "NegativeBinomial" + , [ y; a; none; Helpers.(binop (int 1) Divide (binop (int 1) Plus b)) ] ) | (("neg_binomial_2" | "neg_binomial_2_log") as l), _ -> - raise_s - [%message l " is not supported, consider using neg_binomial instead."] - | f, _ when Operator.of_string_opt f |> Option.is_some -> (fname, args) + raise_s [%message l " is not supported, consider using neg_binomial instead."] + | f, _ when Operator.of_string_opt f |> Option.is_some -> fname, args | _ -> - if Set.mem capitalize_fnames fname then (String.capitalize fname, args) - else raise_s [%message "Not sure how to handle " fname " yet!"] + if Set.mem capitalize_fnames fname + then String.capitalize fname, args + else raise_s [%message "Not sure how to handle " fname " yet!"] +;; let translate_funapps_and_kwrds e = let open Expr.Fixed in - let f ({pattern; _} as expr) = + let f ({ pattern; _ } as expr) = match pattern with | FunApp (StanLib, fname, args) -> - let prefix = - if Utils.is_distribution_name fname then dist_prefix else "" - in - let fname = remove_stan_dist_suffix fname in - let fname, args = map_functions fname args in - {expr with pattern= FunApp (StanLib, prefix ^ fname, args)} + let prefix = if Utils.is_distribution_name fname then dist_prefix else "" in + let fname = remove_stan_dist_suffix fname in + let fname, args = map_functions fname args in + { expr with pattern = FunApp (StanLib, prefix ^ fname, args) } | FunApp (UserDefined, fname, args) -> - { expr with - pattern= FunApp (UserDefined, add_suffix_to_kwrds fname, args) } - | Var s -> {expr with pattern= Var (add_suffix_to_kwrds s)} + { expr with pattern = FunApp (UserDefined, add_suffix_to_kwrds fname, args) } + | Var s -> { expr with pattern = Var (add_suffix_to_kwrds s) } | _ -> expr in rewrite_bottom_up ~f e +;; let%expect_test "nested dist prefixes translated" = let open Expr.Fixed.Pattern in - let e pattern = {Expr.Fixed.pattern; meta= Expr.Typed.Meta.empty} in + let e pattern = { Expr.Fixed.pattern; meta = Expr.Typed.Meta.empty } in let f = FunApp ( Fun_kind.StanLib , "normal_lpdf" - , [FunApp (Fun_kind.StanLib, "normal_lpdf", []) |> e] ) - |> e |> translate_funapps_and_kwrds + , [ FunApp (Fun_kind.StanLib, "normal_lpdf", []) |> e ] ) + |> e + |> translate_funapps_and_kwrds in - print_s [%sexp (f : Expr.Typed.Meta.t Expr.Fixed.t)] ; + print_s [%sexp (f : Expr.Typed.Meta.t Expr.Fixed.t)]; [%expect {| ((pattern @@ -98,54 +131,60 @@ let%expect_test "nested dist prefixes translated" = (((pattern (FunApp StanLib tfd__.Normal ())) (meta ((type_ UInt) (loc ) (adlevel DataOnly))))))) (meta ((type_ UInt) (loc ) (adlevel DataOnly)))) |}] +;; (* temporary until we get rid of these from the MIR *) let rec remove_unused_stmts s = let pattern = match s.Stmt.Fixed.pattern with - | Assignment (_, {Expr.Fixed.pattern= FunApp (CompilerInternal, f, _); _}) + | Assignment (_, { Expr.Fixed.pattern = FunApp (CompilerInternal, f, _); _ }) when Internal_fun.to_string FnConstrain = f - || Internal_fun.to_string FnUnconstrain = f -> - Stmt.Fixed.Pattern.Skip + || Internal_fun.to_string FnUnconstrain = f -> Stmt.Fixed.Pattern.Skip | Decl _ -> Stmt.Fixed.Pattern.Skip - | NRFunApp (CompilerInternal, name, _) - when Internal_fun.to_string FnCheck = name -> - Stmt.Fixed.Pattern.Skip + | NRFunApp (CompilerInternal, name, _) when Internal_fun.to_string FnCheck = name -> + Stmt.Fixed.Pattern.Skip | x -> Stmt.Fixed.Pattern.map Fn.id remove_unused_stmts x in - {s with pattern} + { s with pattern } +;; let rec change_kwrds_stmts s = let open Stmt.Fixed.Pattern in let pattern = match s.Stmt.Fixed.pattern with - | Decl e -> Decl {e with decl_id= add_suffix_to_kwrds e.decl_id} + | Decl e -> Decl { e with decl_id = add_suffix_to_kwrds e.decl_id } | NRFunApp (t, s, e) -> NRFunApp (t, add_suffix_to_kwrds s, e) - | Assignment ((s, t, e1), e2) -> - Assignment ((add_suffix_to_kwrds s, t, e1), e2) - | For e -> For {e with loopvar= add_suffix_to_kwrds e.loopvar} + | Assignment ((s, t, e1), e2) -> Assignment ((add_suffix_to_kwrds s, t, e1), e2) + | For e -> For { e with loopvar = add_suffix_to_kwrds e.loopvar } | x -> map Fn.id change_kwrds_stmts x in - {s with pattern} + { s with pattern } +;; let trans_prog (p : Program.Typed.t) = - let rec map_stmt {Stmt.Fixed.pattern; meta} = - { Stmt.Fixed.pattern= + let rec map_stmt { Stmt.Fixed.pattern; meta } = + { Stmt.Fixed.pattern = Stmt.Fixed.Pattern.map translate_funapps_and_kwrds map_stmt pattern - ; meta } + ; meta + } in - let rename_kwrds (s, e) = (add_suffix_to_kwrds s, e) in - let rename_fdarg (e1, s, e2) = (e1, add_suffix_to_kwrds s, e2) in + let rename_kwrds (s, e) = add_suffix_to_kwrds s, e in + let rename_fdarg (e1, s, e2) = e1, add_suffix_to_kwrds s, e2 in let rename_func (s : 'a Program.fun_def) = { s with - fdname= add_suffix_to_kwrds s.fdname - ; fdargs= List.map ~f:rename_fdarg s.fdargs } + fdname = add_suffix_to_kwrds s.fdname + ; fdargs = List.map ~f:rename_fdarg s.fdargs + } in - Program.map translate_funapps_and_kwrds map_stmt + Program.map + translate_funapps_and_kwrds + map_stmt { p with - output_vars= List.map ~f:rename_kwrds p.output_vars - ; input_vars= List.map ~f:rename_kwrds p.input_vars - ; functions_block= List.map ~f:rename_func p.functions_block } + output_vars = List.map ~f:rename_kwrds p.output_vars + ; input_vars = List.map ~f:rename_kwrds p.input_vars + ; functions_block = List.map ~f:rename_func p.functions_block + } |> Program.map Fn.id change_kwrds_stmts |> Program.map Fn.id remove_unused_stmts |> Program.map_stmts Analysis_and_optimization.Mir_utils.cleanup_empty_stmts +;; diff --git a/test/integration/run_bin_on_args.ml b/test/integration/run_bin_on_args.ml index 52ba25ecb4..8d04d16e07 100644 --- a/test/integration/run_bin_on_args.ml +++ b/test/integration/run_bin_on_args.ml @@ -3,29 +3,30 @@ open Core_kernel let maybe_convert_cmd_to_windows cmd = let pattern = "/install/default/bin/" in let to_windows str = - String.substr_replace_first ~pattern ~with_:"/default.windows/" str - ^ ".exe" + String.substr_replace_first ~pattern ~with_:"/default.windows/" str ^ ".exe" in let path = - String.prefix cmd - (String.substr_index_exn ~pattern cmd + String.length pattern) + String.prefix cmd (String.substr_index_exn ~pattern cmd + String.length pattern) in if Sys.file_exists (to_windows path) then to_windows cmd else cmd +;; let run_capturing_output cmd = let noflags = Array.create ~len:0 "" in let stdout, stdin, stderr = Unix.open_process_full (maybe_convert_cmd_to_windows cmd) noflags in - let chns = [stdout; stderr] in + let chns = [ stdout; stderr ] in let out = List.map ~f:In_channel.input_lines chns in - ignore (Unix.close_process_full (stdout, stdin, stderr)) ; + ignore (Unix.close_process_full (stdout, stdin, stderr)); String.concat ~sep:"\n" (List.concat out) +;; let () = let binary = Sys.argv.(1) in let dirs = Array.(sub Sys.argv ~pos:2 ~len:(length Sys.argv - 2)) in - Array.stable_sort ~compare:String.compare dirs ; + Array.stable_sort ~compare:String.compare dirs; Array.iter dirs ~f:(fun arg -> let cmd = binary ^ " " ^ arg in - Printf.printf " $ %s\n%s\n" cmd (run_capturing_output cmd) ) + Printf.printf " $ %s\n%s\n" cmd (run_capturing_output cmd)) +;; diff --git a/test/unit/Ast_to_Mir_tests.ml b/test/unit/Ast_to_Mir_tests.ml index d21656058b..7ab0a69660 100644 --- a/test/unit/Ast_to_Mir_tests.ml +++ b/test/unit/Ast_to_Mir_tests.ml @@ -13,9 +13,9 @@ let%expect_test "Operator-assign example" = } |} |> trans_prog "" - |> (fun Program.({log_prob; _}) -> log_prob) + |> (fun Program.{ log_prob; _ } -> log_prob) |> Fmt.strf "@[%a@]" (Fmt.list ~sep:Fmt.cut Stmt.Located.pp) - |> print_endline ; + |> print_endline; [%expect {| { @@ -23,9 +23,9 @@ let%expect_test "Operator-assign example" = array[vector[2], 4] x; x[1] = (x[1] ./ r); } |}] +;; -let mir_from_string s = - Frontend_utils.typed_ast_of_string_exn s |> trans_prog "" +let mir_from_string s = Frontend_utils.typed_ast_of_string_exn s |> trans_prog "" let%expect_test "Prefix-Op-Example" = let mir = @@ -39,7 +39,7 @@ let%expect_test "Prefix-Op-Example" = |} in let op = mir.log_prob in - print_s [%sexp (op : Stmt.Located.t list)] ; + print_s [%sexp (op : Stmt.Located.t list)]; (* Perhaps this is producing too many nested lists. XXX*) [%expect {| @@ -68,10 +68,11 @@ let%expect_test "Prefix-Op-Example" = ())) (meta ))))) (meta ))) |}] +;; let%expect_test "read data" = let m = mir_from_string "data { matrix[10, 20] mat[5]; }" in - print_s [%sexp (m.prepare_data : Stmt.Located.t list)] ; + print_s [%sexp (m.prepare_data : Stmt.Located.t list)]; [%expect {| (((pattern @@ -87,10 +88,11 @@ let%expect_test "read data" = ((pattern (Lit Int 5)) (meta ((type_ UInt) (loc ) (adlevel DataOnly))))))))) (meta ))) |}] +;; let%expect_test "read param" = let m = mir_from_string "parameters { matrix[10, 20] mat[5]; }" in - print_s [%sexp (m.log_prob : Stmt.Located.t list)] ; + print_s [%sexp (m.log_prob : Stmt.Located.t list)]; [%expect {| (((pattern @@ -201,12 +203,11 @@ let%expect_test "read param" = (meta ))))) (meta ))))) (meta ))) |}] +;; let%expect_test "gen quant" = - let m = - mir_from_string "generated quantities { matrix[10, 20] mat[5]; }" - in - print_s [%sexp (m.generate_quantities : Stmt.Located.t list)] ; + let m = mir_from_string "generated quantities { matrix[10, 20] mat[5]; }" in + print_s [%sexp (m.generate_quantities : Stmt.Located.t list)]; [%expect {| (((pattern @@ -324,3 +325,4 @@ let%expect_test "gen quant" = (meta ))))) (meta ))))) (meta ))) |}] +;; diff --git a/test/unit/Dataflow_utils.ml b/test/unit/Dataflow_utils.ml index e124d18fd8..b9bedfe126 100644 --- a/test/unit/Dataflow_utils.ml +++ b/test/unit/Dataflow_utils.ml @@ -6,9 +6,8 @@ open Analysis_and_optimization.Dataflow_types let semantic_check_program ast = Option.value_exn - (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Result.ok (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) +;; (***********************************) (* Tests *) @@ -16,7 +15,8 @@ let semantic_check_program ast = let%expect_test "Loop test" = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) @@ -29,19 +29,17 @@ let%expect_test "Loop test" = let statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) - {meta= Location_span.empty; pattern= block}) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) + { meta = Location_span.empty; pattern = block }) in let exits, preds = build_predecessor_graph statement_map in print_s [%sexp - ( statement_map - : ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t ) - Map.Poly.t )] ; - print_s [%sexp (exits : label Set.Poly.t)] ; - print_s [%sexp (preds : (label, label Set.Poly.t) Map.Poly.t)] ; + (statement_map + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t) Map.Poly.t)]; + print_s [%sexp (exits : label Set.Poly.t)]; + print_s [%sexp (preds : (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 @@ -89,10 +87,12 @@ let%expect_test "Loop test" = (1) ((1 (2)) (2 (3)) (3 (4)) (4 (5)) (5 (3))) |}] +;; let%expect_test "Loop passthrough" = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { if (1) { @@ -124,19 +124,21 @@ let%expect_test "Loop passthrough" = let statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) - {meta= Location_span.empty; pattern= block}) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) + { meta = Location_span.empty; pattern = block }) in let exits, _ = build_predecessor_graph statement_map in - print_s [%sexp (exits : label Set.Poly.t)] ; + print_s [%sexp (exits : label Set.Poly.t)]; [%expect {| (1) |}] +;; let example1_program = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { // 1 @@ -172,20 +174,22 @@ let example1_program = in let mir = Ast_to_Mir.trans_prog "" (semantic_check_program ast) in let block = Stmt.Fixed.Pattern.Block mir.log_prob in - Stmt.Fixed.{meta= Location_span.empty; pattern= block} + Stmt.Fixed.{ meta = Location_span.empty; pattern = block } +;; let example1_statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) example1_program) +;; let%expect_test "Statement label map example" = print_s [%sexp - ( Map.Poly.map example1_statement_map ~f:fst - : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t) Map.Poly.t )] ; + (Map.Poly.map example1_statement_map ~f:fst + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t) Map.Poly.t)]; [%expect {| ((1 (Block (2))) (2 (Block (3 4 5))) @@ -282,12 +286,11 @@ let%expect_test "Statement label map example" = (((pattern (Lit Str Fin)) (meta ((type_ UReal) (loc ) (adlevel DataOnly)))))))) |}] +;; let%expect_test "Predecessor graph example" = let exits, preds = build_predecessor_graph example1_statement_map in - print_s - [%sexp - ((exits, preds) : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)] ; + print_s [%sexp (exits, preds : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1) @@ -296,10 +299,11 @@ let%expect_test "Predecessor graph example" = (15 (16)) (16 (14)) (17 (14 15)) (18 (19)) (19 (17)) (20 (21)) (21 (17)) (22 (18 20)))) |}] +;; let%expect_test "Controlflow graph example" = let cf = build_cf_graph example1_statement_map in - print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)] ; + print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 ()) (2 ()) (3 ()) (4 ()) (5 ()) (6 (5)) (7 (5)) (8 (5)) (9 (5 13)) @@ -307,18 +311,22 @@ let%expect_test "Controlflow graph example" = (17 (9 16)) (18 (16 17)) (19 (16 17)) (20 (16 17)) (21 (16 17)) (22 (9 16 19))) |}] +;; let%test "Reconstructed recursive statement" = let stmt = build_recursive_statement - (fun pattern meta -> Stmt.Fixed.{pattern; meta}) - example1_statement_map 1 + (fun pattern meta -> Stmt.Fixed.{ pattern; meta }) + example1_statement_map + 1 in stmt = example1_program +;; let example3_program = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { while (42); @@ -329,24 +337,24 @@ let example3_program = let mir = Ast_to_Mir.trans_prog "" (semantic_check_program ast) in let blocks = Stmt.Fixed.( - Pattern.SList [{pattern= Block mir.log_prob; meta= Location_span.empty}]) + Pattern.SList [ { pattern = Block mir.log_prob; meta = Location_span.empty } ]) in - Stmt.Fixed.{meta= Location_span.empty; pattern= blocks} + Stmt.Fixed.{ meta = Location_span.empty; pattern = blocks } +;; let example3_statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) example3_program) +;; let%expect_test "Statement label map example 3" = print_s [%sexp - ( example3_statement_map - : ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t ) - Map.Poly.t )] ; + (example3_statement_map + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t) Map.Poly.t)]; [%expect {| ((1 @@ -387,31 +395,32 @@ let%expect_test "Statement label map example 3" = (end_loc ((filename string) (line_num 4) (col_num 22) (included_from ()))))))) |}] +;; let%expect_test "Controlflow graph example 3" = let cf = build_cf_graph example3_statement_map in - print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)] ; + print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 ()) (2 ()) (3 ()) (4 ()) (5 (4)) (6 ())) |}] +;; let%expect_test "Predecessor graph example 3" = (* TODO: this is still wrong. The correct answer is - ((6) ((1 ()) (2 (1)) (3 (2)) (4 (3 5)) (5 (4)) (6 (5)))) - Similarly for for-loops. - ) *) + ((6) ((1 ()) (2 (1)) (3 (2)) (4 (3 5)) (5 (4)) (6 (5)))) + Similarly for for-loops. + ) *) let exits, preds = build_predecessor_graph example3_statement_map in - print_s - [%sexp - ((exits, preds) : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)] ; - [%expect - {| + print_s [%sexp (exits, preds : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)]; + [%expect {| ((2) ((1 ()) (2 (3)) (3 (6)) (4 (1 5)) (5 (4)) (6 (4)))) |}] +;; let example4_program = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:6) { @@ -424,24 +433,24 @@ let example4_program = let mir = Ast_to_Mir.trans_prog "" (semantic_check_program ast) in let blocks = Stmt.Fixed.( - Pattern.SList [{pattern= Block mir.log_prob; meta= Location_span.empty}]) + Pattern.SList [ { pattern = Block mir.log_prob; meta = Location_span.empty } ]) in - Stmt.Fixed.{meta= Location_span.empty; pattern= blocks} + Stmt.Fixed.{ meta = Location_span.empty; pattern = blocks } +;; let example4_statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) example4_program) +;; let%expect_test "Statement label map example 4" = print_s [%sexp - ( example4_statement_map - : ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t ) - Map.Poly.t )] ; + (example4_statement_map + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t) Map.Poly.t)]; [%expect {| ((1 @@ -487,33 +496,34 @@ let%expect_test "Statement label map example 4" = (end_loc ((filename string) (line_num 5) (col_num 11) (included_from ()))))))) |}] +;; let%expect_test "Controlflow graph example 4" = let cf = build_cf_graph example4_statement_map in - print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)] ; - [%expect - {| + print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)]; + [%expect {| ((1 ()) (2 ()) (3 ()) (4 ()) (5 (4)) (6 (4)) (7 (4 6))) |}] +;; let%expect_test "Predecessor graph example 4" = let exits, preds = build_predecessor_graph example4_statement_map in (* TODO: this is still wrong. The correct answer is - ( (7) ( (1 ()) (2 (1)) (3 (2)) (4 (3 6)) (5 (4)) (6 (5)) (7 ()) ) ) - or a very conservative approximation - ( (7) ( (1 ()) (2 (1)) (3 (2)) (4 (3 6 7)) (5 (4)) (6 (5)) (7 (6)) ) ) - *) - print_s - [%sexp - ((exits, preds) : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)] ; + ( (7) ( (1 ()) (2 (1)) (3 (2)) (4 (3 6)) (5 (4)) (6 (5)) (7 ()) ) ) + or a very conservative approximation + ( (7) ( (1 ()) (2 (1)) (3 (2)) (4 (3 6 7)) (5 (4)) (6 (5)) (7 (6)) ) ) + *) + print_s [%sexp (exits, preds : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((2) ((1 ()) (2 (3)) (3 (4)) (4 (1 5 6)) (5 (7)) (6 (4)) (7 (6)))) |}] +;; let example5_program = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:6) { @@ -527,24 +537,24 @@ let example5_program = let mir = Ast_to_Mir.trans_prog "" (semantic_check_program ast) in let blocks = Stmt.Fixed.( - Pattern.SList [{pattern= Block mir.log_prob; meta= Location_span.empty}]) + Pattern.SList [ { pattern = Block mir.log_prob; meta = Location_span.empty } ]) in - Stmt.Fixed.{meta= Location_span.empty; pattern= blocks} + Stmt.Fixed.{ meta = Location_span.empty; pattern = blocks } +;; let example5_statement_map = Stmt.Fixed.( build_statement_map - (fun {pattern; _} -> pattern) - (fun {meta; _} -> meta) + (fun { pattern; _ } -> pattern) + (fun { meta; _ } -> meta) example5_program) +;; let%expect_test "Statement label map example 5" = print_s [%sexp - ( example5_statement_map - : ( label - , (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t ) - Map.Poly.t )] ; + (example5_statement_map + : (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * Location_span.t) Map.Poly.t)]; [%expect {| ((1 @@ -595,25 +605,25 @@ let%expect_test "Statement label map example 5" = ((filename string) (line_num 7) (col_num 8) (included_from ()))) (end_loc ((filename string) (line_num 7) (col_num 9) (included_from ()))))))) |}] +;; let%expect_test "Controlflow graph example 5" = let cf = build_cf_graph example5_statement_map in - print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)] ; - [%expect - {| + print_s [%sexp (cf : (label, label Set.Poly.t) Map.Poly.t)]; + [%expect {| ((1 ()) (2 ()) (3 ()) (4 (6)) (5 (4)) (6 (4)) (7 (4)) (8 ())) |}] +;; let%expect_test "Predecessor graph example 5" = let exits, preds = build_predecessor_graph example5_statement_map in (* TODO: this is still very very conservative (e.g. I'd hope for - (8) ((1 ())) (2 (1)) (3 (2)) (4 (3)) (5 (4)) (6 (5)) (7 ()) (8 (6)) - but maybe that's too much to ask for - ) *) - print_s - [%sexp - ((exits, preds) : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)] ; + (8) ((1 ())) (2 (1)) (3 (2)) (4 (3)) (5 (4)) (6 (5)) (7 ()) (8 (6)) + but maybe that's too much to ask for + ) *) + print_s [%sexp (exits, preds : label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((2) ((1 ()) (2 (3)) (3 (8)) (4 (1 5)) (5 (7)) (6 (4)) (7 (6)) (8 (4 6)))) |}] +;; diff --git a/test/unit/Debug_data_generation_tests.ml b/test/unit/Debug_data_generation_tests.ml index e6678c269d..9ab41bb435 100644 --- a/test/unit/Debug_data_generation_tests.ml +++ b/test/unit/Debug_data_generation_tests.ml @@ -5,7 +5,8 @@ open Debug_data_generation let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { int K; int D; @@ -18,11 +19,10 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{\ @@ -33,11 +33,13 @@ let%expect_test "whole program data generation check" = \n 2.56799363086151, 3.3282621325833865, 2.7103944900448411,\ \n 5.2015419032442969, 4.25312636944623]]\ \n}" |}] +;; let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { int x[3, 4]; int y[5, 2, 4]; @@ -50,11 +52,10 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{\ @@ -77,11 +78,13 @@ let%expect_test "whole program data generation check" = \n [4.0434169569431706, 5.2448759493135153, 2.0095894885098069],\ \n [3.8556222147542085, 3.226595023801782, 2.292622453020976]]\ \n}" |}] +;; let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { int K; int D; @@ -95,11 +98,10 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{\ @@ -112,11 +114,13 @@ let%expect_test "whole program data generation check" = \n [3.543689144366625, 6.0288479433993629],\ \n [3.604405750889411, 4.0759938356540726]]\ \n}" |}] +;; let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { corr_matrix[5] d; @@ -134,11 +138,10 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{\ @@ -192,11 +195,13 @@ let%expect_test "whole program data generation check" = \n [1.8124656303877222, 1.8059981193977444, 1.9574266472261275,\ \n 1.3421609989627226]]\ \n}" |}] +;; let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { int N; @@ -226,11 +231,10 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{\ @@ -338,11 +342,13 @@ let%expect_test "whole program data generation check" = \n [1.275024919645803, 0.6402078241901894, 0.],\ \n [1.2238398605980445, 1.0912017876639934, 1.8199435094277936]]]\ \n}" |}] +;; let%expect_test "whole program data generation check" = let open Parse in let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program {| data { int K; // players @@ -356,12 +362,12 @@ let%expect_test "whole program data generation check" = let ast = Option.value_exn (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) in let str = print_data_prog ast in - print_s [%sexp (str : string)] ; + print_s [%sexp (str : string)]; [%expect {| "{ \"K\": 3, \"N\": 1, \"player1\": [2], \"player0\": [1], \"y\": [1]\ \n}" |}] +;; diff --git a/test/unit/Dependence_analysis.ml b/test/unit/Dependence_analysis.ml index 8787ceac52..b75fd8e87b 100644 --- a/test/unit/Dependence_analysis.ml +++ b/test/unit/Dependence_analysis.ml @@ -6,13 +6,13 @@ open Analysis_and_optimization.Dataflow_types let semantic_check_program ast = Option.value_exn - (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Result.ok (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) +;; let example1_program = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { // 1 @@ -47,11 +47,12 @@ let example1_program = |} in Ast_to_Mir.trans_prog "" (semantic_check_program ast) +;; let%expect_test "Dependency graph example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) let deps = log_prob_dependency_graph example1_program in - print_s [%sexp (deps : (label, label Set.Poly.t) Map.Poly.t)] ; + print_s [%sexp (deps : (label, label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 ()) (2 ()) (3 ()) (4 ()) (5 (4)) (6 (4 5)) (7 (4 5)) (8 (4 5)) @@ -61,17 +62,16 @@ let%expect_test "Dependency graph example" = (19 (4 5 9 11 13 14 16 17)) (20 (4 5 9 11 13 14 16 17)) (21 (4 5 9 11 13 14 16 17)) (22 (4 5 9 11 13 14 16 17 19))) |}] +;; let%expect_test "Reaching defns example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) let deps = - Map.Poly.map (log_prob_build_dep_info_map example1_program) - ~f:(fun (_, x) -> + Map.Poly.map (log_prob_build_dep_info_map example1_program) ~f:(fun (_, x) -> ( reaching_defn_lookup x.reaching_defn_entry (VVar "j") - , reaching_defn_lookup x.reaching_defn_exit (VVar "j") ) ) + , reaching_defn_lookup x.reaching_defn_exit (VVar "j") )) in - print_s - [%sexp (deps : (label, label Set.Poly.t * label Set.Poly.t) Map.Poly.t)] ; + print_s [%sexp (deps : (label, label Set.Poly.t * label Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 (() ())) (2 ((9) (9))) (3 (() ())) (4 (() ())) (5 (() ())) (6 (() ())) @@ -80,19 +80,17 @@ let%expect_test "Reaching defns example" = (17 ((9) (9))) (18 ((9) (9))) (19 ((9) (9))) (20 ((9) (9))) (21 ((9) (9))) (22 ((9) (9)))) |}] +;; let%expect_test "Reaching defns example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) let deps = - Map.Poly.map (log_prob_build_dep_info_map example1_program) - ~f:(fun (_, x) -> (x.reaching_defn_entry, x.reaching_defn_exit) ) + Map.Poly.map (log_prob_build_dep_info_map example1_program) ~f:(fun (_, x) -> + x.reaching_defn_entry, x.reaching_defn_exit) in print_s [%sexp - ( deps - : ( label - , reaching_defn Set.Poly.t * reaching_defn Set.Poly.t ) - Map.Poly.t )] ; + (deps : (label, reaching_defn Set.Poly.t * reaching_defn Set.Poly.t) Map.Poly.t)]; [%expect {| ((1 (() ())) (2 ((((VVar i) 4) ((VVar j) 9)) (((VVar i) 4) ((VVar j) 9)))) @@ -115,6 +113,7 @@ let%expect_test "Reaching defns example" = (21 ((((VVar i) 4) ((VVar j) 9)) (((VVar i) 4) ((VVar j) 9)))) (22 ((((VVar i) 4) ((VVar j) 9)) (((VVar i) 4) ((VVar j) 9))))) |}] +;; let%expect_test "Variable dependency example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) @@ -124,14 +123,16 @@ let%expect_test "Variable dependency example" = (Set.Poly.singleton (VVar "j")) 17 in - print_s [%sexp (deps : label Set.Poly.t)] ; + print_s [%sexp (deps : label Set.Poly.t)]; [%expect {| (4 5 9 11 13 14 16) |}] +;; let uninitialized_var_example = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int y) { @@ -178,11 +179,12 @@ let uninitialized_var_example = |} in Ast_to_Mir.trans_prog "" (semantic_check_program ast) +;; let%expect_test "Uninitialized variables example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) let deps = mir_uninitialized_variables uninitialized_var_example in - print_s [%sexp (deps : (Location_span.t * string) Set.Poly.t)] ; + print_s [%sexp (deps : (Location_span.t * string) Set.Poly.t)]; [%expect {| ((((begin_loc @@ -206,3 +208,4 @@ let%expect_test "Uninitialized variables example" = ((filename string) (line_num 42) (col_num 17) (included_from ())))) k)) |}] +;; diff --git a/test/unit/Desugar_test.ml b/test/unit/Desugar_test.ml index 6629fd30a3..77ec6322ab 100644 --- a/test/unit/Desugar_test.ml +++ b/test/unit/Desugar_test.ml @@ -4,10 +4,12 @@ open Analysis_and_optimization let to_mir s = Frontend.Frontend_utils.typed_ast_of_string_exn s |> Frontend.Ast_to_Mir.trans_prog "test prog" +;; -let print_tdata Middle.Program.({prepare_data; _}) = +let print_tdata Middle.Program.{ prepare_data; _ } = Fmt.(strf "@[%a@]@," (list ~sep:cut Middle.Stmt.Located.pp) prepare_data) |> print_endline +;; let%expect_test "matrix array multi indexing " = to_mir @@ -17,13 +19,15 @@ transformed data { matrix[3,4] mat[5]; print(mat[2, arr, 2]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; arr = FnMakeArray__(2, 3, 1); data array[matrix[3, 4], 5] mat; FnPrint__(mat[2, arr, 2]); |}] +;; let%expect_test "matrix array multi indexing " = to_mir @@ -33,7 +37,8 @@ transformed data { matrix[3,4] mat[5]; print(mat[2][arr][2]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; @@ -41,6 +46,7 @@ transformed data { data array[matrix[3, 4], 5] mat; FnPrint__(mat[2][arr[2]]); |}] +;; let%expect_test "matrix array multi indexing " = to_mir @@ -50,7 +56,8 @@ transformed data { matrix[3,4] mat[5]; print(mat[2, arr, arr][2, 2]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; @@ -58,6 +65,7 @@ transformed data { data array[matrix[3, 4], 5] mat; FnPrint__(mat[2, arr[2], arr[2]]); |}] +;; let%expect_test "matrix array multi indexing " = to_mir @@ -67,7 +75,8 @@ transformed data { matrix[3,4] mat[5]; print(mat[3:, 2:3][2, 1]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; @@ -75,6 +84,7 @@ transformed data { data array[matrix[3, 4], 5] mat; FnPrint__(mat[4, 2]); |}] +;; let%expect_test "matrix array multi indexing " = to_mir @@ -84,7 +94,8 @@ transformed data { matrix[3,4] mat[5]; print(mat[:3, 1, :]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; @@ -92,6 +103,7 @@ transformed data { data array[matrix[3, 4], 5] mat; FnPrint__(mat[1:3, 1]); |}] +;; let%expect_test "matrix array multi indexing " = to_mir @@ -105,7 +117,8 @@ transformed data { print(mat[2, :, arr][2, 1]); print(mat[:, 2, arr][2, 1]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data array[int, 3] arr; @@ -116,6 +129,7 @@ transformed data { FnPrint__(mat[2, arr[1], 1]); FnPrint__(mat[2, 2, arr[1]]); FnPrint__(mat[2, 2, arr[1]]); |}] +;; let%expect_test "intertwined with partial evaluator" = to_mir {| @@ -123,7 +137,9 @@ transformed data { vector[3] x; print(log(1-x[:])[:]); } |} - |> Partial_evaluator.eval_prog |> print_tdata ; + |> Partial_evaluator.eval_prog + |> print_tdata; [%expect {| data vector[3] x; FnPrint__(log1m(x)); |}] +;; diff --git a/test/unit/Factor_graph.ml b/test/unit/Factor_graph.ml index 1e119defde..a36ab62206 100644 --- a/test/unit/Factor_graph.ml +++ b/test/unit/Factor_graph.ml @@ -5,13 +5,13 @@ open Analysis_and_optimization.Dataflow_types let semantic_check_program ast = Option.value_exn - (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Result.ok (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) +;; let reject_example = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { real x; @@ -49,11 +49,12 @@ let reject_example = |} in Ast_to_Mir.trans_prog "" (semantic_check_program ast) +;; let%expect_test "Factor graph reject example" = (*let deps = snd (build_predecessor_graph example1_statement_map) in*) let deps = prog_factor_graph reject_example in - print_s [%sexp (deps : factor_graph)] ; + print_s [%sexp (deps : factor_graph)]; [%expect {| ((factor_map @@ -70,10 +71,12 @@ let%expect_test "Factor graph reject example" = ()))) (var_map ())) |}] +;; let complex_example = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { real a; @@ -101,10 +104,11 @@ let complex_example = |} in Ast_to_Mir.trans_prog "" (semantic_check_program ast) +;; let%expect_test "Factor graph complex example" = let deps = prog_factor_graph complex_example in - print_s [%sexp (deps : factor_graph)] ; + print_s [%sexp (deps : factor_graph)]; [%expect {| ((factor_map @@ -326,10 +330,12 @@ let%expect_test "Factor graph complex example" = (meta ((type_ UReal) (loc ) (adlevel AutoDiffable))))) 21)))))) |}] +;; let complex_example = let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| data { real x; @@ -357,11 +363,11 @@ let complex_example = |} in Ast_to_Mir.trans_prog "" (semantic_check_program ast) +;; let%expect_test "Priors complex example" = let priors = list_priors complex_example in - print_s - [%sexp (priors : (vexpr, (factor * label) Set.Poly.t option) Map.Poly.t)] ; + print_s [%sexp (priors : (vexpr, (factor * label) Set.Poly.t option) Map.Poly.t)]; [%expect {| (((VVar a) @@ -423,3 +429,4 @@ let%expect_test "Priors complex example" = 13)))) ((VVar c) (())) ((VVar d) (())) ((VVar e) (())) ((VVar f) (()))) |}] +;; diff --git a/test/unit/Optimize.ml b/test/unit/Optimize.ml index 45e1237c92..dc06c4c735 100644 --- a/test/unit/Optimize.ml +++ b/test/unit/Optimize.ml @@ -7,14 +7,14 @@ open Analysis_and_optimization.Mir_utils let semantic_check_program ast = Option.value_exn - (Result.ok - (Semantic_check.semantic_check_program - (Option.value_exn (Result.ok ast)))) + (Result.ok (Semantic_check.semantic_check_program (Option.value_exn (Result.ok ast)))) +;; let%expect_test "map_rec_stmt_loc" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { print(24); @@ -30,12 +30,12 @@ let%expect_test "map_rec_stmt_loc" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let f = function - | Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [s]) -> - Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [s; s]) + | Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [ s ]) -> + Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [ s; s ]) | x -> x in let mir = Program.map Fn.id (map_rec_stmt_loc f) mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -54,11 +54,13 @@ let%expect_test "map_rec_stmt_loc" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "map_rec_state_stmt_loc" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { print(24); @@ -74,18 +76,17 @@ let%expect_test "map_rec_state_stmt_loc" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let f i = function - | Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [s]) -> - Stmt.Fixed.Pattern. - (NRFunApp (CompilerInternal, "FnPrint__", [s; s]), i + 1) - | x -> (x, i) + | Stmt.Fixed.Pattern.NRFunApp (CompilerInternal, "FnPrint__", [ s ]) -> + Stmt.Fixed.Pattern.(NRFunApp (CompilerInternal, "FnPrint__", [ s; s ]), i + 1) + | x -> x, i in let mir_stmt, num = (map_rec_state_stmt_loc f 0) - Stmt.Fixed.{pattern= SList mir.log_prob; meta= Location_span.empty} + Stmt.Fixed.{ pattern = SList mir.log_prob; meta = Location_span.empty } in - let mir = {mir with log_prob= [mir_stmt]} in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; - print_endline (string_of_int num) ; + let mir = { mir with log_prob = [ mir_stmt ] } in + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; + print_endline (string_of_int num); [%expect {| log_prob { @@ -108,11 +109,13 @@ let%expect_test "map_rec_state_stmt_loc" = 3 |}] +;; let%expect_test "inline functions" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { void f(int x, matrix y) { @@ -132,7 +135,7 @@ let%expect_test "inline functions" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -178,11 +181,13 @@ let%expect_test "inline functions" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline functions 2" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { void f() { @@ -199,7 +204,7 @@ let%expect_test "inline functions 2" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -235,11 +240,13 @@ let%expect_test "inline functions 2" = if(inline_sym7__) break; } } |}] +;; let%expect_test "list collapsing" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { void f(int x, matrix y) { @@ -260,7 +267,7 @@ let%expect_test "list collapsing" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in let mir = list_collapsing mir in - print_s [%sexp (mir : Middle.Program.Typed.t)] ; + print_s [%sexp (mir : Middle.Program.Typed.t)]; [%expect {| ((functions_block @@ -440,11 +447,13 @@ let%expect_test "list collapsing" = (meta )))) (transform_inits ()) (output_vars ()) (prog_name "") (prog_path "")) |}] +;; let%expect_test "do not inline recursive functions" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { real g(int z); @@ -460,7 +469,7 @@ let%expect_test "do not inline recursive functions" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -489,11 +498,13 @@ let%expect_test "do not inline recursive functions" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function in for loop" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -513,7 +524,7 @@ let%expect_test "inline function in for loop" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -576,13 +587,15 @@ let%expect_test "inline function in for loop" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; (* TODO: check test results from here *) let%expect_test "inline function in for loop 2" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -602,7 +615,7 @@ let%expect_test "inline function in for loop 2" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -685,11 +698,13 @@ let%expect_test "inline function in for loop 2" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function in while loop" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -709,7 +724,7 @@ let%expect_test "inline function in while loop" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -761,11 +776,13 @@ let%expect_test "inline function in while loop" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function in if then else" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -785,7 +802,7 @@ let%expect_test "inline function in if then else" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -829,11 +846,13 @@ let%expect_test "inline function in if then else" = } |}] +;; let%expect_test "inline function in ternary if " = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -857,7 +876,7 @@ let%expect_test "inline function in ternary if " = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -926,11 +945,13 @@ let%expect_test "inline function in ternary if " = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function multiple returns " = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -949,7 +970,7 @@ let%expect_test "inline function multiple returns " = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -993,11 +1014,13 @@ let%expect_test "inline function multiple returns " = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function indices " = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -1014,7 +1037,7 @@ let%expect_test "inline function indices " = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -1060,11 +1083,13 @@ let%expect_test "inline function indices " = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function and " = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -1081,7 +1106,7 @@ let%expect_test "inline function and " = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -1128,11 +1153,13 @@ let%expect_test "inline function and " = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "inline function or " = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int f(int z) { @@ -1148,7 +1175,7 @@ let%expect_test "inline function or " = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = function_inlining mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -1195,11 +1222,13 @@ let%expect_test "inline function or " = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "unroll nested loop" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) for (j in 3:4) @@ -1210,7 +1239,7 @@ let%expect_test "unroll nested loop" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = static_loop_unrolling mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1238,11 +1267,13 @@ let%expect_test "unroll nested loop" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "unroll nested loop 2" = let _ = Gensym.reset_danger_use_cautiously () in let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) for (j in i:4) @@ -1254,7 +1285,7 @@ let%expect_test "unroll nested loop 2" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = static_loop_unrolling mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1437,11 +1468,13 @@ let%expect_test "unroll nested loop 2" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "unroll nested loop 3" = let _ = Gensym.reset_danger_use_cautiously () in let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) for (j in i:4) @@ -1453,7 +1486,7 @@ let%expect_test "unroll nested loop 3" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = static_loop_unrolling mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1534,11 +1567,13 @@ let%expect_test "unroll nested loop 3" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "unroll nested loop with break" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) for (j in 3:4) { @@ -1551,7 +1586,7 @@ let%expect_test "unroll nested loop with break" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = static_loop_unrolling mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1575,11 +1610,13 @@ let%expect_test "unroll nested loop with break" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "constant propagation" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| transformed data { int i; @@ -1597,7 +1634,7 @@ let%expect_test "constant propagation" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = constant_propagation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| prepare_data { @@ -1619,11 +1656,13 @@ let%expect_test "constant propagation" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "constant propagation, local scope" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| transformed data { int i; @@ -1644,7 +1683,7 @@ let%expect_test "constant propagation, local scope" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = constant_propagation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| prepare_data { @@ -1669,11 +1708,13 @@ let%expect_test "constant propagation, local scope" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "constant propagation, model block local scope" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1693,7 +1734,7 @@ let%expect_test "constant propagation, model block local scope" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = constant_propagation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1720,11 +1761,13 @@ let%expect_test "constant propagation, model block local scope" = generated_quantities int i; //int generated_quantities int j; //int } |}] +;; let%expect_test "expression propagation" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| transformed data { int i; @@ -1741,7 +1784,7 @@ let%expect_test "expression propagation" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = expression_propagation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| prepare_data { @@ -1762,11 +1805,13 @@ let%expect_test "expression propagation" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "copy propagation" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1783,7 +1828,7 @@ let%expect_test "copy propagation" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = copy_propagation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1803,11 +1848,13 @@ let%expect_test "copy propagation" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "dead code elimination" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| transformed data { int i[2]; @@ -1826,7 +1873,7 @@ let%expect_test "dead code elimination" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| prepare_data { @@ -1849,11 +1896,13 @@ let%expect_test "dead code elimination" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "dead code elimination decl" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1870,7 +1919,7 @@ let%expect_test "dead code elimination decl" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1887,11 +1936,13 @@ let%expect_test "dead code elimination decl" = FnPrint__(i); } } |}] +;; let%expect_test "dead code elimination, for loop" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1903,7 +1954,7 @@ let%expect_test "dead code elimination, for loop" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1917,11 +1968,13 @@ let%expect_test "dead code elimination, for loop" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "dead code elimination, while loop" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1937,7 +1990,7 @@ let%expect_test "dead code elimination, while loop" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -1952,11 +2005,13 @@ let%expect_test "dead code elimination, while loop" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "dead code elimination, if then" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -1982,7 +2037,7 @@ let%expect_test "dead code elimination, if then" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2002,11 +2057,13 @@ let%expect_test "dead code elimination, if then" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "dead code elimination, nested" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int i; @@ -2020,7 +2077,7 @@ let%expect_test "dead code elimination, nested" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2034,11 +2091,13 @@ let%expect_test "dead code elimination, nested" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "partial evaluation" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { if (1 > 2) { @@ -2053,7 +2112,7 @@ let%expect_test "partial evaluation" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = partial_evaluation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2071,11 +2130,13 @@ let%expect_test "partial evaluation" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "try partially evaluate" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { real x; @@ -2090,7 +2151,7 @@ let%expect_test "try partially evaluate" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = partial_evaluation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2108,11 +2169,13 @@ let%expect_test "try partially evaluate" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "partially evaluate with equality check" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { vector[2] x; @@ -2125,7 +2188,7 @@ let%expect_test "partially evaluate with equality check" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = partial_evaluation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2141,11 +2204,13 @@ let%expect_test "partially evaluate with equality check" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "partially evaluate functions" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { matrix[3, 2] x_matrix; @@ -2282,7 +2347,7 @@ model { let mir = Ast_to_Mir.trans_prog "" ast in let mir = constant_propagation mir in let mir = partial_evaluation mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2458,11 +2523,13 @@ model { parameters real theta_u; //real parameters real phi_u; //real } |}] +;; let%expect_test "lazy code motion" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { print({3.0}); @@ -2475,7 +2542,7 @@ let%expect_test "lazy code motion" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2494,11 +2561,13 @@ let%expect_test "lazy code motion" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 2" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { for (i in 1:2) @@ -2510,7 +2579,7 @@ let%expect_test "lazy code motion, 2" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2528,11 +2597,13 @@ let%expect_test "lazy code motion, 2" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 3" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { print(3); @@ -2545,7 +2616,7 @@ let%expect_test "lazy code motion, 3" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2565,11 +2636,13 @@ let%expect_test "lazy code motion, 3" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 4" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int b; @@ -2595,7 +2668,7 @@ let%expect_test "lazy code motion, 4" = let mir = list_collapsing mir in (* TODO: make sure that these temporaries do not get assigned level DataOnly unless appropriate *) - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2632,11 +2705,13 @@ let%expect_test "lazy code motion, 4" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 5" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int b; @@ -2660,7 +2735,7 @@ let%expect_test "lazy code motion, 5" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2701,11 +2776,13 @@ let%expect_test "lazy code motion, 5" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 6" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int x; @@ -2720,7 +2797,7 @@ let%expect_test "lazy code motion, 6" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2740,11 +2817,13 @@ let%expect_test "lazy code motion, 6" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 7" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int a; @@ -2777,7 +2856,7 @@ let%expect_test "lazy code motion, 7" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2817,11 +2896,13 @@ let%expect_test "lazy code motion, 7" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 8, _lp functions not optimized" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| functions { int foo_lp(int x) { target += 1; return 24; } @@ -2839,7 +2920,7 @@ let%expect_test "lazy code motion, 8, _lp functions not optimized" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| functions { @@ -2875,11 +2956,13 @@ let%expect_test "lazy code motion, 8, _lp functions not optimized" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 9" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int x; @@ -2891,7 +2974,7 @@ let%expect_test "lazy code motion, 9" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2908,11 +2991,13 @@ let%expect_test "lazy code motion, 9" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 10" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int x; @@ -2927,7 +3012,7 @@ let%expect_test "lazy code motion, 10" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2947,11 +3032,13 @@ let%expect_test "lazy code motion, 10" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 11" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { { @@ -2969,7 +3056,7 @@ let%expect_test "lazy code motion, 11" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -2992,11 +3079,13 @@ let%expect_test "lazy code motion, 11" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 12" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { int x; @@ -3011,7 +3100,7 @@ let%expect_test "lazy code motion, 12" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -3031,11 +3120,13 @@ let%expect_test "lazy code motion, 12" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "lazy code motion, 13" = let _ = Gensym.reset_danger_use_cautiously () in let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { real temp; @@ -3058,7 +3149,7 @@ let%expect_test "lazy code motion, 13" = let mir = one_step_loop_unrolling mir in let mir = lazy_code_motion mir in let mir = list_collapsing mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -3101,12 +3192,15 @@ let%expect_test "lazy code motion, 13" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; -let%expect_test "cool example: expression propagation + partial evaluation + \ - lazy code motion + dead code elimination" = - Gensym.reset_danger_use_cautiously () ; +let%expect_test "cool example: expression propagation + partial evaluation + lazy code \ + motion + dead code elimination" + = + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| model { real x; @@ -3127,7 +3221,7 @@ let%expect_test "cool example: expression propagation + partial evaluation + \ let mir = lazy_code_motion mir in let mir = list_collapsing mir in let mir = dead_code_elimination mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -3156,12 +3250,12 @@ let%expect_test "cool example: expression propagation + partial evaluation + \ if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "block fixing" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program - {| + Parse.parse_string Parser.Incremental.program {| model { } |} @@ -3170,21 +3264,25 @@ let%expect_test "block fixing" = let mir = Ast_to_Mir.trans_prog "" ast in let mir = { mir with - Middle.Program.log_prob= + Middle.Program.log_prob = [ Stmt.Fixed. - { pattern= + { pattern = IfElse ( Expr.Helpers.zero - , { pattern= + , { pattern = While ( Expr.Helpers.zero - , {pattern= SList []; meta= Location_span.empty} ) - ; meta= Location_span.empty } + , { pattern = SList []; meta = Location_span.empty } ) + ; meta = Location_span.empty + } , None ) - ; meta= Location_span.empty } ] } + ; meta = Location_span.empty + } + ] + } in let mir = block_fixing mir in - print_s [%sexp (mir : Program.Typed.t)] ; + print_s [%sexp (mir : Program.Typed.t)]; [%expect {| ((functions_block ()) (input_vars ()) (prepare_data ()) @@ -3226,11 +3324,13 @@ let%expect_test "block fixing" = ((pattern (Return ())) (meta )) ())) (meta )))) (transform_inits ()) (output_vars ()) (prog_name "") (prog_path "")) |}] +;; let%expect_test "one-step loop unrolling" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| transformed data { int x; @@ -3243,7 +3343,7 @@ let%expect_test "one-step loop unrolling" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = one_step_loop_unrolling mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| prepare_data { @@ -3281,11 +3381,13 @@ let%expect_test "one-step loop unrolling" = if(PNot__(emit_transformed_parameters__ || emit_generated_quantities__)) return; if(PNot__(emit_generated_quantities__)) return; } |}] +;; let%expect_test "adlevel_optimization" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { real w; @@ -3313,7 +3415,7 @@ let%expect_test "adlevel_optimization" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = optimize_ad_levels mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -3355,11 +3457,13 @@ let%expect_test "adlevel_optimization" = output_vars { parameters real w; //real } |}] +;; let%expect_test "adlevel_optimization expressions" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { real w; @@ -3387,7 +3491,7 @@ let%expect_test "adlevel_optimization expressions" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = optimize_ad_levels mir in - print_s [%sexp (mir.log_prob : Stmt.Located.t list)] ; + print_s [%sexp (mir.log_prob : Stmt.Located.t list)]; [%expect {| (((pattern @@ -3483,11 +3587,13 @@ let%expect_test "adlevel_optimization expressions" = (meta ((type_ UReal) (loc ) (adlevel DataOnly))))))) (meta ))))) (meta ))) |}] +;; let%expect_test "adlevel_optimization 2" = - Gensym.reset_danger_use_cautiously () ; + Gensym.reset_danger_use_cautiously (); let ast = - Parse.parse_string Parser.Incremental.program + Parse.parse_string + Parser.Incremental.program {| parameters { real w; @@ -3516,7 +3622,7 @@ let%expect_test "adlevel_optimization 2" = let ast = semantic_check_program ast in let mir = Ast_to_Mir.trans_prog "" ast in let mir = optimize_ad_levels mir in - Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline ; + Fmt.strf "@[%a@]" Program.Typed.pp mir |> print_endline; [%expect {| log_prob { @@ -3563,3 +3669,4 @@ let%expect_test "adlevel_optimization 2" = parameters real w; //real transformed_parameters real w_trans; //real } |}] +;; diff --git a/test/unit/Parse_tests.ml b/test/unit/Parse_tests.ml index 9302d81484..56eedf531c 100644 --- a/test/unit/Parse_tests.ml +++ b/test/unit/Parse_tests.ml @@ -7,12 +7,11 @@ let render_syntax_error = Fmt.to_to_string Errors.pp_syntax_error (* TESTS *) let%expect_test "parse conditional" = let ast = - parse_string Parser.Incremental.program - "model { if (1 < 2) { print(\"hi\");}}" + parse_string Parser.Incremental.program "model { if (1 < 2) { print(\"hi\");}}" |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -31,16 +30,18 @@ let%expect_test "parse conditional" = ())) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse dangling else problem" = let ast = - parse_string Parser.Incremental.program - "model { if (1 < 2) print(\"I'm sorry\"); if (2 < 3) print(\", Dave, \ - \"); else print(\"I'm afraid I can't do that.\");}" + parse_string + Parser.Incremental.program + "model { if (1 < 2) print(\"I'm sorry\"); if (2 < 3) print(\", Dave, \"); else \ + print(\"I'm afraid I can't do that.\");}" |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -66,6 +67,7 @@ let%expect_test "parse dangling else problem" = (smeta ((loc ))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse minus unary" = let ast = @@ -73,7 +75,7 @@ let%expect_test "parse minus unary" = |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -98,6 +100,7 @@ let%expect_test "parse minus unary" = (emeta ((loc ))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse unary over binary" = let ast = @@ -105,7 +108,7 @@ let%expect_test "parse unary over binary" = |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -138,15 +141,17 @@ let%expect_test "parse unary over binary" = (is_global false))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse indices, two different colons" = let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program "model { matrix[5, 5] x; print(x[2 - 3 ? 3 : 4 : 2]); }" |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -182,16 +187,17 @@ let%expect_test "parse indices, two different colons" = (emeta ((loc )))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse operator precedence" = let ast = - parse_string Parser.Incremental.program - "model { \ - print({a,b?c:d||e&&f==g!=h<=i=k>l+m-n*o/p%q.*s./t\\r^u[v]'}); }" + parse_string + Parser.Incremental.program + "model { print({a,b?c:d||e&&f==g!=h<=i=k>l+m-n*o/p%q.*s./t\\r^u[v]'}); }" |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -355,10 +361,12 @@ let%expect_test "parse operator precedence" = (emeta ((loc )))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse crazy truncation example" = let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program "\n\ \ model {\n\ \ real T[1,1] = {{42.0}};\n\ @@ -369,7 +377,7 @@ let%expect_test "parse crazy truncation example" = |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -421,10 +429,12 @@ let%expect_test "parse crazy truncation example" = (emeta ((loc )))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; let%expect_test "parse nested loop" = let ast = - parse_string Parser.Incremental.program + parse_string + Parser.Incremental.program " model {\n\ \ for (i in 1:2)\n\ \ for (j in 3:4)\n\ @@ -434,7 +444,7 @@ let%expect_test "parse nested loop" = |> Result.map_error ~f:render_syntax_error |> Result.ok_or_failwith in - print_s [%sexp (ast : Ast.untyped_program)] ; + print_s [%sexp (ast : Ast.untyped_program)]; [%expect {| ((functionblock ()) (datablock ()) (transformeddatablock ()) @@ -455,3 +465,4 @@ let%expect_test "parse nested loop" = (smeta ((loc ))))))) (smeta ((loc ))))))) (generatedquantitiesblock ())) |}] +;; diff --git a/test/unit/Pedantic_analysis.ml b/test/unit/Pedantic_analysis.ml index 79f0dbde3b..98c8f91ae3 100644 --- a/test/unit/Pedantic_analysis.ml +++ b/test/unit/Pedantic_analysis.ml @@ -3,12 +3,14 @@ open Frontend open Analysis_and_optimization.Pedantic_analysis let build_program prog = - Ast_to_Mir.trans_prog "" + Ast_to_Mir.trans_prog + "" (Option.value_exn (Result.ok (Semantic_check.semantic_check_program (Option.value_exn (Result.ok (Parse.parse_string Parser.Incremental.program prog)))))) +;; let sigma_example = {| @@ -30,9 +32,10 @@ let sigma_example = x ~ normal (0, z); } |} +;; let%expect_test "Unbounded sigma warning" = - print_warn_pedantic (build_program sigma_example) ; + print_warn_pedantic (build_program sigma_example); [%expect {| Warning: @@ -60,6 +63,7 @@ let%expect_test "Unbounded sigma warning" = A normal distribution is given value -1 as a scale parameter (argument 2), but a scale parameter is not strictly positive. |}] +;; let uniform_example = {| @@ -77,9 +81,10 @@ let uniform_example = d ~ uniform(0, 1); } |} +;; let%expect_test "Uniform warning" = - print_warn_pedantic (build_program uniform_example) ; + print_warn_pedantic (build_program uniform_example); [%expect {| Warning: @@ -112,6 +117,7 @@ let%expect_test "Uniform warning" = constraints; for example, instead of giving an elasticity parameter a uniform(0,1) distribution, try normal(0.5,0.5). |}] +;; let unscaled_example = {| @@ -129,9 +135,10 @@ let unscaled_example = z = -1000 + 0.00001; } |} +;; let%expect_test "Unscaled warning" = - print_warn_pedantic (build_program unscaled_example) ; + print_warn_pedantic (build_program unscaled_example); [%expect {| Warning at 'string', line 11, column 21 to column 26: @@ -140,6 +147,7 @@ let%expect_test "Unscaled warning" = Warning at 'string', line 11, column 28 to column 33: Argument 10000 suggests there may be parameters that are not unit scale; consider rescaling with a multiplier (see manual section 22.12). |}] +;; let multi_twiddle_example = {| @@ -153,9 +161,10 @@ let multi_twiddle_example = x ~ normal(y, 1); } |} +;; let%expect_test "Multi twiddle warning" = - print_warn_pedantic (build_program multi_twiddle_example) ; + print_warn_pedantic (build_program multi_twiddle_example); [%expect {| Warning: @@ -165,6 +174,7 @@ let%expect_test "Multi twiddle warning" = Warning at 'string', line 7, column 10 to column 27: The parameter x is on the left-hand side of more than one twiddle statement. |}] +;; let hard_constrained_example = {| @@ -179,9 +189,10 @@ let hard_constrained_example = model { } |} +;; let%expect_test "Hard constraint warning" = - print_warn_pedantic (build_program hard_constrained_example) ; + print_warn_pedantic (build_program hard_constrained_example); [%expect {| Warning: @@ -222,6 +233,7 @@ let%expect_test "Hard constraint warning" = soft constraints rather than hard constraints; for example, instead of constraining an elasticity parameter to fall between 0, and 1, leave it unconstrained and give it a normal(0.5,0.5) prior distribution. |}] +;; let unused_param_example = {| @@ -243,9 +255,10 @@ let unused_param_example = real g = d; } |} +;; let%expect_test "Unused param warning" = - print_warn_pedantic (build_program unused_param_example) ; + print_warn_pedantic (build_program unused_param_example); [%expect {| Warning: @@ -258,6 +271,7 @@ let%expect_test "Unused param warning" = The parameter e was declared but was not used in the density calculation. Warning: The parameter f was declared but was not used in the density calculation. |}] +;; let param_dependant_cf_example = {| @@ -281,9 +295,10 @@ let param_dependant_cf_example = } } |} +;; let%expect_test "Parameter dependent control flow warning" = - print_warn_pedantic (build_program param_dependant_cf_example) ; + print_warn_pedantic (build_program param_dependant_cf_example); [%expect {| Warning at 'string', line 9, column 10 to line 13, column 11: @@ -292,6 +307,7 @@ let%expect_test "Parameter dependent control flow warning" = A control flow statement depends on parameter(s): a. Warning at 'string', line 17, column 10 to line 19, column 11: A control flow statement depends on parameter(s): a. |}] +;; let non_one_priors_example = {| @@ -313,9 +329,10 @@ let non_one_priors_example = x ~ normal(c, d); } |} +;; let%expect_test "Non-one priors no warning" = - print_warn_pedantic (build_program non_one_priors_example) ; + print_warn_pedantic (build_program non_one_priors_example); [%expect {| Warning at 'string', line 15, column 24 to column 25: @@ -324,6 +341,7 @@ let%expect_test "Non-one priors no warning" = Warning at 'string', line 17, column 24 to column 25: A normal distribution is given parameter d as a scale parameter (argument 2), but d was not constrained to be strictly positive. |}] +;; let non_one_priors_example2 = {| @@ -351,9 +369,10 @@ let non_one_priors_example2 = f ~ normal(e, 1); } |} +;; let%expect_test "Non-one priors warning" = - print_warn_pedantic (build_program non_one_priors_example2) ; + print_warn_pedantic (build_program non_one_priors_example2); [%expect {| Warning: @@ -371,6 +390,7 @@ let%expect_test "Non-one priors warning" = Warning at 'string', line 22, column 10 to column 27: The parameter f is on the left-hand side of more than one twiddle statement. |}] +;; let gamma_args_example = {| @@ -389,9 +409,10 @@ let gamma_args_example = d ~ gamma(0.4, 0.6); } |} +;; let%expect_test "Gamma args warning" = - print_warn_pedantic (build_program gamma_args_example) ; + print_warn_pedantic (build_program gamma_args_example); [%expect {| Warning: @@ -424,6 +445,7 @@ let%expect_test "Gamma args warning" = A inv_gamma distribution is given parameter b as a scale parameter (argument 2), but b was not constrained to be strictly positive. |}] +;; let dist_bounds_example = {| @@ -440,9 +462,10 @@ let dist_bounds_example = d ~ lognormal(2, 2); } |} +;; let%expect_test "Dist bounds warning" = - print_warn_pedantic (build_program dist_bounds_example) ; + print_warn_pedantic (build_program dist_bounds_example); [%expect {| Warning at 'string', line 9, column 10 to column 11: @@ -452,6 +475,7 @@ let%expect_test "Dist bounds warning" = Parameter c is given a lognormal distribution, which has strictly positive support, but c was not constrained to be strictly positive. |}] +;; let dist_examples = {| @@ -593,11 +617,12 @@ model { cov ~ inv_wishart(pos_p, cov); } |} +;; (* Distribution warnings should appear only on alternating lines, since the program lines go incorrect,correct,incorrect,correct,etc.*) let%expect_test "Dist warnings" = - print_warn_pedantic (build_program dist_examples) ; + print_warn_pedantic (build_program dist_examples); [%expect {| Warning: @@ -958,6 +983,7 @@ let%expect_test "Dist warnings" = A inv_wishart distribution is given parameter mat as a scale matrix (argument 2), but mat was not constrained to be covariance. |}] +;; let fundef_cf_example = {| @@ -981,9 +1007,10 @@ model { x ~ normal(0, func(sigma)); } |} +;; let%expect_test "Function body parameter-dependent control flow" = - print_warn_pedantic (build_program fundef_cf_example) ; + print_warn_pedantic (build_program fundef_cf_example); [%expect {| Warning: @@ -993,3 +1020,4 @@ let%expect_test "Function body parameter-dependent control flow" = 'string', line 19, column 21 to column 26, the value of b depends on parameter(s): sigma. |}] +;; diff --git a/test/unit/Semantic_check_tests.ml b/test/unit/Semantic_check_tests.ml index 130843be90..187ff6d0ba 100644 --- a/test/unit/Semantic_check_tests.ml +++ b/test/unit/Semantic_check_tests.ml @@ -12,7 +12,7 @@ transformed data { |} |> typed_ast_of_string_exn |> Fmt.strf "@[%a@]" Pretty_printing.pp_program - |> print_endline ; + |> print_endline; [%expect {| transformed data { @@ -20,3 +20,4 @@ transformed data { matrix[3, 4] mat[5]; print(mat[indices, : , indices][2, 1, 1]); } |}] +;; diff --git a/test/unit/Stan_math_code_gen_tests.ml b/test/unit/Stan_math_code_gen_tests.ml index a6000afd0a..ef21221267 100644 --- a/test/unit/Stan_math_code_gen_tests.ml +++ b/test/unit/Stan_math_code_gen_tests.ml @@ -5,22 +5,23 @@ open Fmt open Stan_math_code_gen let%expect_test "udf" = - let with_no_loc stmt = - Stmt.Fixed.{pattern= stmt; meta= Locations.no_span_num} - in - let w e = Expr.{Fixed.pattern= e; meta= Typed.Meta.empty} in + let with_no_loc stmt = Stmt.Fixed.{ pattern = stmt; meta = Locations.no_span_num } in + let w e = Expr.{ Fixed.pattern = e; meta = Typed.Meta.empty } in let pp_fun_def_w_rs a b = pp_fun_def a b String.Set.empty in - { fdrt= None - ; fdname= "sars" - ; fdargs= [(DataOnly, "x", UMatrix); (AutoDiffable, "y", URowVector)] - ; fdbody= + { fdrt = None + ; fdname = "sars" + ; fdargs = [ DataOnly, "x", UMatrix; AutoDiffable, "y", URowVector ] + ; fdbody = Stmt.Fixed.Pattern.Return - (Some - (w @@ FunApp (StanLib, "add", [w @@ Var "x"; w @@ Lit (Int, "1")]))) - |> with_no_loc |> List.return |> Stmt.Fixed.Pattern.Block |> with_no_loc - ; fdloc= Location_span.empty } + (Some (w @@ FunApp (StanLib, "add", [ w @@ Var "x"; w @@ Lit (Int, "1") ]))) + |> with_no_loc + |> List.return + |> Stmt.Fixed.Pattern.Block + |> with_no_loc + ; fdloc = Location_span.empty + } |> strf "@[%a" pp_fun_def_w_rs - |> print_endline ; + |> print_endline; [%expect {| template @@ -52,3 +53,4 @@ let%expect_test "udf" = return sars(x, y, pstream__); } }; |}] +;;