Skip to content

Commit 0f352a6

Browse files
WardBrianWardBrian
committed
Add more linear functions, cleanup
1 parent 00195e5 commit 0f352a6

3 files changed

Lines changed: 45 additions & 85 deletions

File tree

‎src/analysis_and_optimization/Pedantic_analysis.ml‎

Lines changed: 30 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -69,50 +69,55 @@ let list_multi_tildes (mir : Program.Typed.t) :
6969
(** These functions are linear if all of their arguments are *)
7070
let linear_fnames =
7171
Operator.([Plus; PPlus; Minus; PMinus] |> List.map ~f:to_string)
72-
@ [ "block"; "col"; "cols"; "row"; "rows"; "diagonal"; "head"; "tail"
73-
; "negative_infinity"; "not_a_number"; "rep_matrix"; "positive_infinity"
74-
; "segment"; "sum"; "to_vector" ]
72+
@ [ "add"; "append_block"; "append_row"; "append_col"; "block"; "col"; "cols"
73+
; "row"; "rows"; "diagonal"; "head"; "tail"; "minus"; "negative_infinity"
74+
; "not_a_number"; "rep_matrix"; "rep_vector"; "rep_row_vector"
75+
; "positive_infinity"; "segment"; "subtract"; "sum"; "to_vector"
76+
; "to_row_vector"; "to_matrix"; "to_array_1d"; "to_array_2d"; "transpose"
77+
; "Plus__"; "PPlus__"; "PMinus__"; "Minus__"; "PNot__"; "Transpose__" ]
7578
|> String.Set.of_list
7679

7780
let list_possible_nonlinear (mir : Program.Typed.t) : Location_span.t Set.Poly.t
7881
=
7982
(* Collect statements of the form "target += Dist(param, ...)" *)
8083
let rec is_linear_function allow_var name (args : 'a Expr.Fixed.t list) =
81-
if Set.mem linear_fnames name then
82-
List.for_all ~f:(fun {pattern; _} -> is_linear allow_var pattern) args
83-
(* A function of all constants is fine *)
84-
else if List.for_all ~f:(fun {pattern; _} -> is_linear false pattern) args
85-
then true
86-
else
87-
match (name, args) with
88-
(* We require at least one of these operands to be a constant *)
89-
| ("Times__" | "Divide__" | "IntDivide__"), [a; b] ->
90-
(is_linear allow_var a.pattern && is_linear false b.pattern)
91-
|| (is_linear false a.pattern && is_linear allow_var b.pattern)
92-
(* Transpose of transpose is fine *)
93-
| ( "Transpose__"
94-
, [{pattern= FunApp (StanLib ("Transpose__", _, _), [a]); _}] ) ->
95-
is_linear allow_var a.pattern
96-
| _ -> false
97-
and is_linear allow_var = function
84+
match (name, args) with
85+
| _, _ when Set.mem linear_fnames name ->
86+
List.for_all ~f:(is_linear allow_var) args
87+
| _, _ when List.for_all ~f:(is_linear false) args ->
88+
(* A function of all constants is fine *) true
89+
(* We require at least one of these operands to be a constant *)
90+
| ("Times__" | "Divide__" | "IntDivide__"), [a; b] ->
91+
(is_linear allow_var a && is_linear false b)
92+
|| (is_linear false a && is_linear allow_var b)
93+
(* Partial evaluation can create fmas where the user wrote Times*)
94+
| "fma", [a; b; c] ->
95+
is_linear allow_var c
96+
&& ( (is_linear allow_var a && is_linear false b)
97+
|| (is_linear false a && is_linear allow_var b) )
98+
| _ -> false
99+
and is_linear allow_var Expr.Fixed.{pattern; _} =
100+
match pattern with
98101
| Expr.Fixed.Pattern.Var _ -> allow_var
99102
| Lit _ -> true
100-
| Indexed (Expr.Fixed.{pattern; _}, _) -> is_linear allow_var pattern
103+
| Indexed (e, _) | Promotion (e, _, _) -> is_linear allow_var e
104+
| TernaryIf (e1, e2, e3) ->
105+
is_linear allow_var e1 && is_linear allow_var e2
106+
&& is_linear allow_var e3
101107
| FunApp (StanLib (name, _, _), args) ->
102108
is_linear_function allow_var name args
103109
| FunApp (CompilerInternal (FnMakeArray | FnMakeRowVec), args) ->
104-
List.for_all ~f:(fun {pattern; _} -> is_linear allow_var pattern) args
110+
List.for_all ~f:(is_linear allow_var) args
105111
| _ -> false in
106112
let collect_transformed_tilde_stmt (stmt : Stmt.Located.t) :
107113
Location_span.t Set.Poly.t =
108114
match stmt.pattern with
109115
| Stmt.Fixed.Pattern.TargetPE
110116
{ pattern=
111117
Expr.Fixed.Pattern.FunApp
112-
( (StanLib (_, FnLpdf _, _) | UserDefined (_, FnLpdf _))
113-
, {pattern; _} :: _ )
118+
((StanLib (_, FnLpdf _, _) | UserDefined (_, FnLpdf _)), e :: _)
114119
; _ }
115-
when not (is_linear true pattern) ->
120+
when not (is_linear true e) ->
116121
Set.Poly.singleton stmt.meta
117122
| _ -> Set.Poly.empty in
118123
let tildes =

‎test/integration/cli-args/warn-pedantic/jacobian_false_alarm.stan‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ data {
55
parameters {
66
vector[3] v;
77
array[4,5] real a;
8-
matrix[6,6] m;
8+
cholesky_factor_corr[6] m;
99
real y;
1010
real z;
1111
}
@@ -20,22 +20,22 @@ model {
2020
y ~ normal(0,1);
2121

2222
// fun
23-
m ~ lkj_corr(2.0);
24-
(m + m) ~ lkj_corr(2.0);
25-
(m - m) ~ lkj_corr(2.0);
26-
(v + v) ~ multi_normal(v,m);
27-
(v - v) ~ multi_normal(v,m);
28-
block(m,1,1,1,1) ~ lkj_corr(2.0);
23+
m ~ lkj_corr_cholesky(2.0);
24+
(m + m) ~ lkj_corr_cholesky(2.0);
25+
(m - m) ~ lkj_corr_cholesky(2.0);
26+
(v + v) ~ normal(0,1);
27+
(v - v) ~ normal(0,1);
28+
block(m,1,1,1,1) ~ lkj_corr_cholesky(2.0);
2929
col(m,1) ~ normal(0,1);
3030
cols(m) ~ normal(0,1);
3131
row(m,1) ~ normal(0,1);
3232
rows(m) ~ normal(0,1);
33-
diagonal(m) ~ multi_normal(v,m);
34-
head(v,2) ~ multi_normal(v,m);
33+
diagonal(m) ~ normal(0,1);
34+
head(v,2) ~ normal(0,1);
3535
negative_infinity() ~ normal(0,1);
3636
not_a_number() ~ normal(0,1);
37-
rep_matrix(1,3,3) ~ lkj_corr(2.0);
38-
(v')' ~ multi_normal(v,m);
37+
rep_matrix(1,3,3) ~ lkj_corr_cholesky(2.0);
38+
(v')' ~ normal(0,1);
3939
positive_infinity() ~ normal(0,1);
4040
segment(v,2,4) ~ normal(0,1);
4141
sum(v) ~ normal(0,1);
@@ -64,6 +64,7 @@ model {
6464

6565
// literals
6666
[1] ~ normal(0,1);
67+
[y + y] ~ normal(0,1);
6768
to_vector({1}) ~ normal(0,1);
6869

6970
}

‎test/integration/cli-args/warn-pedantic/stanc.expected‎

Lines changed: 3 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -400,56 +400,10 @@ Warning: Parameter f has constraints that don't make sense. The lower bound
400400
Warning: Parameter e has constraints that don't make sense. The lower bound
401401
should be strictly less than the upper bound.
402402
$ ../../../../../install/default/bin/stanc --warn-pedantic jacobian_false_alarm.stan
403-
Warning in 'jacobian_false_alarm.stan', line 58, column 2: Left-hand side of
404-
sampling statement (~) may contain a non-linear transform of a parameter
405-
or local variable. If it does, you need to include a target += statement
406-
with the log absolute determinant of the Jacobian of the transform.
407-
Warning in 'jacobian_false_alarm.stan', line 38, column 25: A multi_normal
408-
distribution is given parameter m as a covariance matrix (argument 2),
409-
but m was not constrained to be covariance.
410-
Warning in 'jacobian_false_alarm.stan', line 37, column 2: Left-hand side of
411-
sampling statement (~) may contain a non-linear transform of a parameter
412-
or local variable. If it does, you need to include a target += statement
413-
with the log absolute determinant of the Jacobian of the transform.
414-
Warning in 'jacobian_false_alarm.stan', line 37, column 2: It is suggested to
415-
reparameterize your model to replace lkj_corr with lkj_corr_cholesky, the
416-
Cholesky factor variant. lkj_corr tends to run slower, consume more
417-
memory, and has higher risk of numerical errors.
418-
Warning in 'jacobian_false_alarm.stan', line 34, column 29: A multi_normal
419-
distribution is given parameter m as a covariance matrix (argument 2),
420-
but m was not constrained to be covariance.
421-
Warning in 'jacobian_false_alarm.stan', line 33, column 31: A multi_normal
422-
distribution is given parameter m as a covariance matrix (argument 2),
423-
but m was not constrained to be covariance.
424-
Warning in 'jacobian_false_alarm.stan', line 28, column 2: It is suggested to
425-
reparameterize your model to replace lkj_corr with lkj_corr_cholesky, the
426-
Cholesky factor variant. lkj_corr tends to run slower, consume more
427-
memory, and has higher risk of numerical errors.
428-
Warning in 'jacobian_false_alarm.stan', line 27, column 27: A multi_normal
429-
distribution is given parameter m as a covariance matrix (argument 2),
430-
but m was not constrained to be covariance.
431-
Warning in 'jacobian_false_alarm.stan', line 26, column 27: A multi_normal
432-
distribution is given parameter m as a covariance matrix (argument 2),
433-
but m was not constrained to be covariance.
434-
Warning in 'jacobian_false_alarm.stan', line 25, column 2: It is suggested to
435-
reparameterize your model to replace lkj_corr with lkj_corr_cholesky, the
436-
Cholesky factor variant. lkj_corr tends to run slower, consume more
437-
memory, and has higher risk of numerical errors.
438-
Warning in 'jacobian_false_alarm.stan', line 24, column 2: It is suggested to
439-
reparameterize your model to replace lkj_corr with lkj_corr_cholesky, the
440-
Cholesky factor variant. lkj_corr tends to run slower, consume more
441-
memory, and has higher risk of numerical errors.
442-
Warning in 'jacobian_false_alarm.stan', line 23, column 2: It is suggested to
443-
reparameterize your model to replace lkj_corr with lkj_corr_cholesky, the
444-
Cholesky factor variant. lkj_corr tends to run slower, consume more
445-
memory, and has higher risk of numerical errors.
446-
Warning in 'jacobian_false_alarm.stan', line 23, column 2: Parameter m is
447-
given a lkj_corr distribution, which has correlation support, but m was
448-
not constrained to be correlation.
449403
Warning: The parameter z has 7 priors.
450-
Warning: The parameter y has 8 priors.
451-
Warning: The parameter v has 9 priors.
452-
Warning: The parameter m has 16 priors.
404+
Warning: The parameter y has 9 priors.
405+
Warning: The parameter v has 8 priors.
406+
Warning: The parameter m has 12 priors.
453407
Warning: The parameter a has 2 priors.
454408
$ ../../../../../install/default/bin/stanc --warn-pedantic jacobian_warning1.stan
455409
Warning in 'jacobian_warning1.stan', line 7, column 2: Left-hand side of

0 commit comments

Comments
 (0)