Merge pull request #800 from stan-dev/jacobian-plus-equal

Basic docs for jacobian +=

Author: bob-carpenter

src/reference-manual/deprecations.qmd

@@ -52,6 +52,14 @@ would use Cholesky factors rather than full correlation matrix types.
 
 *Scheduled Removal*: Stan 3.0 or later.
 
+## Use of `_lp` functions in `transformed parameters`
+
+*Deprecated*: Using [functions that end in `_lp`](user-functions.qmd#log-probability-access-in-functions)
+in the `transformed parameters` block.
+
+*Replacement*: Use `_jacobian` functions and the `jacobian +=` statement instead. These allow for change-of-variable
+adjustments which can be conditionally enabled by Stan's algorithms.
+
 ## New Keywords
 
 *Deprecated*: The following identifiers will become reserved in the language in the specified version.

src/reference-manual/grammar.txt

@@ -24,9 +24,6 @@
 
 <identifier> ::= IDENTIFIER
                | TRUNCATE
-               | <future_keyword>
-
-<future_keyword> ::= JACOBIAN
 
 <decl_identifier> ::= <identifier>
                     | <reserved_word>
@@ -86,6 +83,7 @@
 
 <unsized_type> ::= ARRAY <unsized_dims> <basic_type>
                  | ARRAY <unsized_dims> <unsized_tuple_type>
+                 | <basic_type> <unsized_dims>
                  | <basic_type>
                  | <unsized_tuple_type>
 
@@ -115,7 +113,9 @@
                                   <id_and_optional_assignment(rhs,
                                   <decl_identifier_after_comma>)>)*
 
-<decl(type_rule, rhs)> ::= <higher_type(type_rule)>
+<decl(type_rule, rhs)> ::= type_rule <decl_identifier> LBRACK <expression>
+                           (COMMA <expression>)* RBRACK
+                         | <higher_type(type_rule)>
                            <id_and_optional_assignment(rhs,
                            <decl_identifier>)>
                            [<remaining_declarations(rhs)>] SEMICOLON
@@ -272,6 +272,7 @@
                      | <expression> TILDE <identifier> LPAREN [<expression>
                        (COMMA <expression>)*] RPAREN [<truncation>] SEMICOLON
                      | TARGET PLUSASSIGN <expression> SEMICOLON
+                     | JACOBIAN PLUSASSIGN <expression> SEMICOLON
                      | BREAK SEMICOLON
                      | CONTINUE SEMICOLON
                      | PRINT LPAREN <printables> RPAREN SEMICOLON

src/reference-manual/statements.qmd

@@ -15,7 +15,7 @@ are used --- if they do not, the behavior is undefined.
 The basis of Stan's execution is the evaluation of a log probability
 function (specifically, a log probability density function) for a given
 set of (real-valued) parameters. Log probability functions can be
-constructed by using distribution statements and log probability increment 
+constructed by using distribution statements and log probability increment
 statements.  Statements may be grouped
 into sequences and into for-each loops.  In addition, Stan allows
 local variables to be declared in blocks and also allows an empty
@@ -248,7 +248,8 @@ of the posterior up to an additive constant.  Data and transformed
 data are fixed before the log density is evaluated.  The total log
 probability is initialized to zero.  Next, any log Jacobian
 adjustments accrued by the variable constraints are added to the log
-density (the Jacobian adjustment may be skipped for optimization).
+density (the Jacobian adjustment may be skipped for maximum likelihood estimation
+via optimization).
 Distribution statements and log probability increment statements may add to the log
 density in the model block.  A log probability increment statement
 directly increments the log density with the value of an expression as
@@ -361,14 +362,51 @@ including arrays of vectors and matrices.   For container arguments,
 their sum will be added to the total log density.
 
 
+## Increment log density with a change of variables adjustment
+
+
+A variant of the `target +=` statement described above is the
+`jacobian +=` statement. This can be used in the transformed parameters block
+or in functions ending with `_jacobian` to mimic the log Jacobian
+adjustments accrued by built-in variable transforms.
+
+Similarly to those implemented for the built-in transforms, these Jacobian adjustment
+may be skipped for maximum likelihood estimation via optimization.
+
+For example, here is a program which recreates the existing
+[`<upper=x>` transform](transforms.qmd#upper-bounded-scalar) on real numbers:
+
+```stan
+functions {
+  real upper_bound_jacobian(real x, real ub) {
+    jacobian += x;
+    return ub - exp(x);
+  }
+}
+data {
+  real ub;
+}
+parameters {
+  real b_raw;
+}
+transformed parameters {
+  real b = upper_bound_jacobian(b_raw, ub);
+}
+model {
+  // use b as if it was declared `real<upper=ub> b;` in parameters
+  // e.g.
+  // b ~ lognormal(0, 1);
+}
+```
+
 ### Accessing the log density {-}
 
-To access accumulated log density up to the current execution point,
+To access the accumulated log density up to the current execution point,
 the function `target()` may be used.
 
 ## Sampling statements {#sampling-statements.section}
 
-The term "sampling statement" has been replaced with 
+The term "sampling statement" has been replaced with
 [distribution statement](#distribution-statements.section).
 
 ## Distribution statements {#distribution-statements.section}
@@ -432,7 +470,7 @@ terms in the model block.  Equivalently, each $\sim$ statement
 corresponds to a multiplicative factor in the unnormalized posterior
 density.
 
-Distribution statements (`~`) accept only built-in or user-defined 
+Distribution statements (`~`) accept only built-in or user-defined
 distributions on the
 right side. The left side of a distribution statement may be data,
 parameter, or a complex expression, but the evaluated type needs to
@@ -452,8 +490,8 @@ target += normal_lpdf(sigma | 0, 1);
 ```
 
 Stan models can mix distribution statements and log probability
-increment statements. Although statistical models 
-are usually defined with distributions in the literature, 
+increment statements. Although statistical models
+are usually defined with distributions in the literature,
 there are several scenarios in which we may want to code the log
 likelihood or parts of it directly, for example, due to computational
 efficiency (e.g. censored data model) or coding language limitations
@@ -485,7 +523,7 @@ target += dist_lpmf(y | theta1, ..., thetaN);
 
 This will be well formed if and only if `dist_lpdf(y | theta1,
   ..., thetaN)` or `dist_lpmf(y | theta1, ..., thetaN)` is a
-well-formed expression of type `real`. User defined distributions 
+well-formed expression of type `real`. User defined distributions
 can be defined in functions block by using function names ending with `_lpdf`.
 
 
@@ -881,7 +919,7 @@ The equivalent code for a vectorized truncation depends on which of the
 variables are non-scalars (arrays, vectors, etc.):
 
 1. If the variate `y` is the only non-scalar, the result is the same as
-   described in the above sections, but the `lcdf`/`lccdf` calculation is 
+   described in the above sections, but the `lcdf`/`lccdf` calculation is
    multiplied by `size(y)`.
 
 2. If the other arguments to the distribution are non-scalars, then the

src/reference-manual/syntax.qmd

@@ -300,6 +300,7 @@ The raw output is available [here](https://raw.githubusercontent.com/stan-dev/do
                      | <expression> TILDE <identifier> LPAREN [<expression>
                        (COMMA <expression>)*] RPAREN [<truncation>] SEMICOLON
                      | TARGET PLUSASSIGN <expression> SEMICOLON
+                     | JACOBIAN PLUSASSIGN <expression> SEMICOLON
                      | BREAK SEMICOLON
                      | CONTINUE SEMICOLON
                      | PRINT LPAREN <printables> RPAREN SEMICOLON
@@ -476,6 +477,9 @@ of other user defined functions which end in `_lp`.
 Sampling statements (using `~`) can only be used in the `model` block or in the
 bodies of user-defined functions which end in `_lp`.
 
+`jacobian +=` statements can only be used inside of the `transformed parameters` block
+or in functions that end with `_jacobian`.
+
 ### Probability function naming {-}
 
 A probability function literal must have one of the following

src/reference-manual/user-functions.qmd

@@ -95,7 +95,7 @@ arguments to produce an expression, which has a value when executed.
 ### Functions as statements {-}
 
 Functions with void return types may be applied to arguments and used
-as [statements.qmd](statements). 
+as [statements.qmd](statements).
 These act like distribution statements or print
 statements.  Such uses are only appropriate for functions that act
 through side effects, such as incrementing the log probability
@@ -161,7 +161,11 @@ used in place of parameterized distributions on the right side of
 Functions of certain types are restricted on scope of usage.
 Functions whose names end in `_lp` assume access to the log
 probability accumulator and are only available in the transformed
-parameter and model blocks.
+parameters and model blocks.
+
+Functions whose name end in `_jacobian` assume access to the log
+probability accumulator may only be used within the transformed parameters
+block.
 
 Functions whose names end in `_rng`
 assume access to the random number generator and may only be used
@@ -293,8 +297,8 @@ a function elsewhere results in a compile-time error.
 
 ### Log probability access in functions {-}
 
-Functions that include 
-[statements.qmd#distribution-statements.section](distribution statements) or 
+Functions that include
+[statements.qmd#distribution-statements.section](distribution statements) or
 [statements.qmd#increment-log-prob.section](log probability increment statements)
 must have a name that ends in `_lp`.
 Attempts to use distribution statements or increment log probability

src/stan-users-guide/reparameterization.qmd

@@ -441,10 +441,10 @@ parameters {
 transformed parameters {
   real<lower=0> y;
   y = 1 / y_inv;  // change variables
+  jacobian += -2 * log(y_inv); // Jacobian adjustment
 }
 model {
   y ~ gamma(2,4);
-  target +=  -2 * log(y_inv);  //  Jacobian adjustment;
 }
 ```
 
@@ -476,7 +476,7 @@ transformed parameters {
   matrix[K, K] J;   // Jacobian matrix of transform
   // ... compute v as a function of u ...
   // ... compute J[m, n] = d.v[m] / d.u[n] ...
-  target += log(abs(determinant(J)));
+  jacobian += log(abs(determinant(J)));
   // ...
 }
 model {
@@ -505,7 +505,7 @@ transformed parameters {
   vector[K] J_diag;  // diagonals of Jacobian matrix
   // ...
   // ... compute J[k, k] = d.v[k] / d.u[k] ...
-  target += sum(log(J_diag));
+  jacobian += sum(log(J_diag));
   // ...
 }
 ```
@@ -596,10 +596,10 @@ parameters {
 }
 transformed parameters {
   vector[N] alpha = L + exp(alpha_raw);
+  jacobian += sum(alpha_raw); // log Jacobian
   // ...
 }
 model {
-  target += sum(alpha_raw);  // log Jacobian
   // ...
 }
 ```

src/stan-users-guide/user-functions.qmd

@@ -287,6 +287,41 @@ transformed parameters {
 }
 ```
 
+## Functions implementing change-of-variable adjustments
+
+Functions whose names end in `_jacobian` can use the
+`jacobian +=` statement. This can be used to implement a custom
+change of variables for arbitrary parameters.
+
+For example, this function recreates the built-in
+`<upper=x>` transform on real numbers:
+```stan
+real upper_bound_jacobian(real x, real ub) {
+  jacobian += x;
+  return ub - exp(x);
+}
+```
+
+It can be used as a replacement for `real<lower=ub>` as follows:
+
+```stan
+functions {
+  // upper_bound_jacobian as above
+}
+data {
+  real ub;
+}
+parameters {
+  real b_raw;
+}
+transformed parameters {
+  real b = upper_bound_jacobian(b_raw, ub);
+}
+model {
+  b ~ lognormal(0, 1);
+  // ...
+}
+```
 
 ## Functions acting as random number generators
 

src/stan-users-guide/using-stanc.qmd

@@ -647,7 +647,9 @@ Warning:
     Left-hand side of distribution statement (~) may contain a non-linear
     transform of a parameter or local variable. If it does, you need
     to include a target += statement with the log absolute determinant
-    of the Jacobian of the transform.
+    of the Jacobian of the transform. You could also consider defining
+    a transformed parameter and using jacobian += in the transformed
+    parameters block.
 ```
 
 ### Pedantic mode limitations {-}