Allow users to declare terms as Jacobian

Background:
Currently, when you run optimize, by default Stan does not include the Jacobian term for the constrained parameters (so that with uniform priors, you get the maximum-likelihood estimate). However, when people do constraint transforms manually, Stan does not recognize those and the Jacobian terms are included in all optimize calls. This has potential to produce some very surprising behaviour for some users.

This came up in the context of discussions by @jonah and others on allowing the use of the result of optimize to form a Laplace approximation and draw samples from there (similarly to how we generate samples from ADVI), which would make the optimize method more useful, especially in early workflow stages and so the necessity for it to work for most models without modificiaton could potentially increase.

Idea:
So I wanted to resurface the idea which was already suggested a few times by @Bob_Carpenter (e.g., here and here) to have something like a jacobian += x statement. In a simple implementation, it would just translate to:

if(include_jacobian) {
  target += x;
}

Where include_jacobian is true always except for optimize calls (the switch already exists for the built-in Jacobians, so should be easy to use elsewhere).

Bob originally suggested that such a statement would also be permissible in transformed parameters, and while I see the appeal, I don’t think that’s in any way necessary for it to be useful.

A possible positive side-effect would be more readable code: I can use jacobian += to document that a term is a Jacobian and that the programmer was aware that a Jacobian adjustment was needed for a sampling statement or similar. (this actually speaks against allowing the statement in transfromed parameters).

More speculatively, we could store the jacobian separately and allow users to inspect it, which could have some debugging potential (there are recurring topics where people wonder why the lp__ of a fitted model differs from the sum of the terms in the model block - e.g. here, so having the jacobian accessible could clarify a bit).

Also pedantic mode could actually detect whether a transformed value that is on the left hand side of a sampling statement has an associated jacobian += statement to reduce false positives of the jacobian warning (but that would require some guessing on side of the parser, so maybe not completely practical).

It could also make sense to be more explicit in the naming (because it is the log-jacobian).

This is currently not an official suggestion, I also don’t directly need this for anything, but I promised to post it after the discussion at the Gathering, so I am bringing this up to see what people think. If there’s interest we would then write a design-doc.

7 Likes