Implement more consistent tracking of logp components via LogJacobianAccumulator

[penelopeysm]

As discussed in today's meeting, we want to separate logjac from logprior so that we can more correctly calculate log probs in unlinked / model space.

This PR therefore introduces:

• LogJacobianAccumulator, which accumulates the log-jacobian of the link transform, if the variable is linked. That is, logjac is zero for unlinked variables and nonzero for linked variables (if the link transform is not identity).

Warning

This represents a sign change with respect to how the Jacobian was previously treated. See below.

The main reason for this is because I primarily think of linking as an 'active' choice, that is to say, logjac should be nonzero when linked (i.e. a correction is needed to change the value of logp from the default, which is unlinked). This means that "including logjac" gets you logprior_internal and logjoint_internal, whereas "excluding logjac" means logprior and logjoint.

• Previously, the logjac term was stuffed into LogPriorAccumulator. Now, LogPriorAccumulator only calculates the prior of the unlinked variables i.e. $\log(p(\mathbf{x}))$ below.

• New logp-extracting functions to get the following quantities:

  • getlogjac: $\log(|\det \mathbf{J}|)$
  • getlogprior: $\log(p(\mathbf{x}))$
  • getlogprior_internal: $\log(q(\mathbf{y}))$
  • getlogjoint: $\log(p(\mathbf{x})) + \text{log-likelihood}$
  • getlogjoint_internal: $\log(q(\mathbf{y})) + \text{log-likelihood}$

Warning

Prior to this PR, getlogprior would return $\log(q(\mathbf{y}))$ and getlogjoint would return $\log(q(\mathbf{y})) + \text{log-likelihood}$. @mhauru and I agree that this behaviour is slightly unprincipled because the log prior (and log joint) should have a well-defined meaning independent of whether the VarInfo has been linked or not (which is an implementation detail that most users should not see). See TuringLang/Turing.jl#2617 (comment)

---

Convention: Link transform

For a constrained variable $\mathbf{x} \sim p(\cdot)$ that is transformed to an unconstrained variable $\mathbf{y} \sim q(\cdot)$, we have that

$$\log(q(\mathbf{y})) = \log(p(\mathbf{x})) - \log(|\det \mathbf{J}|)$$

where

$$\mathbf{J} = \begin{pmatrix} (\partial y_1/\partial x_1) & \cdots \\ \ldots & \ddots \\ \end{pmatrix}$$

In Bijectors this is calculated as

b = bijector(p)
y, logjac = with_logabsdet_jacobian(b, x)

In this PR, logjac always refers to this link transform Jacobian. The new accumulator always tracks this quantity.

There are one or two occasions where instead the invlink transform Jacobian is calculated. This PR renames all such occurrences to inv_logjac:

b_inv = inverse(bijector(p))
x, inv_logjac = with_logabsdet_jacobian(b_inv, y)

Once this is merged I will open a PR to the docs page to explain this. The notation used here follows directly on from what I previously wrote at https://turinglang.org/docs/developers/transforms/dynamicppl/.

Closes #992.

[github-actions]

Benchmark Report for Commit 049b3d3

Computer Information

Julia Version 1.11.6
Commit 9615af0f269 (2025-07-09 12:58 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                 13.4 |                 1.3 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                732.7 |                38.8 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                519.1 |                43.9 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1029.5 |                33.2 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               7270.1 |                27.4 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1537.2 |                26.1 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1102.1 |                12.1 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               6657.2 |                52.8 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |                988.5 |                 8.9 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              80565.3 |                48.0 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8602.6 |                 9.7 |
|               Dynamic |        10 |    mooncake |             typed |   true |                143.9 |                22.5 |
|              Submodel |         1 |    mooncake |             typed |   true |                 16.9 |                25.9 |
|                   LDA |        12 | reversediff |             typed |   true |               1159.2 |                 1.9 |

[codecov]

Codecov Report

Attention: Patch coverage is 85.26316% with 14 lines in your changes missing coverage. Please review.

Project coverage is 82.07%. Comparing base (e60eab0) to head (049b3d3).
Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/simple_varinfo.jl	14.28%	6 Missing ⚠️
src/logdensityfunction.jl	20.00%	4 Missing ⚠️
src/default_accumulators.jl	88.88%	3 Missing ⚠️
src/abstract_varinfo.jl	95.45%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking     #998      +/-   ##
============================================
- Coverage     82.11%   82.07%   -0.04%     
============================================
  Files            38       38              
  Lines          4031     4067      +36     
============================================
+ Hits           3310     3338      +28     
- Misses          721      729       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

DynamicPPL.jl documentation for PR #998 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR998/

[penelopeysm]

I think this should be reviewable. Truthfully, most of the changes are either accumulator boilerplate (I'll open a separate issue to try to reduce this), or finicky bits about where and what sign to accumulate logjac.

I don't recommend checking every one of them in review but I will say that the test suite did catch a lot of errors and force me to fix them carefully, so we can have some confidence that this does behave as expected.

The main design decision is probably the sign convention to use. This is covered in the top comment.

[penelopeysm]

I can't say I fully appreciate the rationale for this slightly counterintuitive behaviour. I would be quite happy to remove it. At least changing the variable name hopefully makes it clearer what's going on (I spent ages trying to find out what was going wrong with the accumulators...).

[penelopeysm]

This line happily fixes the following bug:

julia> using DynamicPPL

julia> svi = SimpleVarInfo()
SimpleVarInfo(NamedTuple(), 0.0)

julia> svi_linked = DynamicPPL.settrans!!(svi, true)
Transformed SimpleVarInfo(NamedTuple(), 0.0)

julia> istrans(svi_linked)
true

julia> istrans(DynamicPPL.ThreadSafeVarInfo(svi_linked))
false

After this PR, the last line will return true.

[mhauru]

Looks very nice, some comments and questions.

[mhauru]

A bit unsure about the name logJ, which isn't snake_case. Do you have a reason to prefer it over logjac?

[penelopeysm]

It matches the maths notation, and is consistent with Bijectors.jl:

pysm@ati:~/ppl/bi (main) $ rg logJ | wc -l
      43

I don't know to what extent snake_case matters for things that are mathematical variables.

If all the float-accumulators are unified (and presumably the field will be called logp or similar), will this still be a problem?

[mhauru]

In the unification I'm introducing a function called logp, so the field name would remain. Happy to be consistent with Bijectors.

[penelopeysm]

I do realise that in the rest of DynamicPPL we use logjac though, so for library-internal-consistency's sake we should change it. I'll do a big sed later.

[mhauru]

Is there a chance that some mix-up of using a different invlink transform than what was originally used for linking would cause the logjac to actually be non-zero? Or would that always imply that quite a serious error has been made and we have no need to have well-defined behaviour?

[penelopeysm]

OK, true, I guess somebody could do something horrific and use different StaticTransformations, in which case we should sum the logjac terms and add / subtract them as necessary.

[mhauru]

Why is the distribution new?

[penelopeysm]

Because of a typo 😅

[mhauru]

Same question as above.

[mhauru]

yay!