Basic docs for jacobian +=

Author: WardBrian

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

@@ -361,6 +361,38 @@ 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 optimization.
+
+For example, here is a program which re-creates 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);
+}
+```
+
 ### Accessing the log density {-}
 
 To access accumulated log density up to the current execution point,

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/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,32 @@ 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, here is a program which re-creates the built-in
+`<upper=x>` transform 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);
+}
+```
 
 ## 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 {-}