Implement more consistent tracking of logp components via `LogJacobianAccumulator` #998

penelopeysm · 2025-07-21T17:08:17Z

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 · 2025-07-21T17:19:29Z

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 · 2025-07-21T17:25:32Z

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 · 2025-07-22T00:12:50Z

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.

src/transforming.jl

penelopeysm · 2025-07-22T00:29:20Z

test/simple_varinfo.jl

-            retval, vi_linked_result = DynamicPPL.evaluate!!(model, deepcopy(vi_linked))
+            # NOTE: Evaluating a linked VarInfo, **specifically when the transformation
+            # is static**, will result in an invlinked VarInfo. This is because of
+            # `maybe_invlink_before_eval!`, which only invlinks if the transformation
+            # is static. (src/abstract_varinfo.jl)
+            retval, vi_unlinked_again = DynamicPPL.evaluate!!(model, deepcopy(vi_linked))


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 · 2025-07-22T00:32:22Z

src/simple_varinfo.jl

 istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, vn::VarName) = istrans(vi.varinfo, vn)
+istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}) = istrans(vi.varinfo)


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.

docs/make.jl

mhauru · 2025-07-23T12:42:09Z

src/abstract_varinfo.jl

@@ -306,6 +361,19 @@ function acclogprior!!(vi::AbstractVarInfo, logp)
    return map_accumulator!!(acc -> acc + LogPriorAccumulator(logp), vi, Val(:LogPrior))
 end

+"""
+    acclogjac!!(vi::AbstractVarInfo, logJ)


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

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?

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

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 · 2025-07-23T12:45:03Z

src/abstract_varinfo.jl

+    vi_new = unflatten(vi, x)
+    # Reset logjac to 0.
+    if hasacc(vi_new, Val(:LogJacobian))
+        vi_new = map_accumulator!!(zero, vi_new, Val(:LogJacobian))


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?

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 · 2025-07-23T12:46:25Z

src/accumulators.jl

 - `Base.copy(acc::T)`

+In these functions:
+- `val` is the new value of the random variable sampled from a new distribution (always


Why is the distribution new?

Because of a typo 😅

src/context_implementations.jl

mhauru · 2025-07-23T13:13:01Z

src/simple_varinfo.jl

-        vi_new = acclogprior!!(vi_new, logjac)
+    # logjac should be zero for an unlinked VarInfo.
+    if hasacc(vi_new, Val(:LogJacobian))
+        vi_new = map_accumulator!!(zero, vi_new, Val(:LogJacobian))


Same question as above.

src/transforming.jl

mhauru

yay!

penelopeysm added 2 commits July 21, 2025 17:36

logjac accumulator

a29b953

Fix tests

d6c9cfa

github-actions bot assigned penelopeysm Jul 21, 2025

Fix a whole bunch of stuff

e671a56

penelopeysm changed the title ~~LogJacobianAccumulator~~ LogJacobianAccumulator + more consistent tracking of logp components Jul 21, 2025

penelopeysm changed the title ~~LogJacobianAccumulator + more consistent tracking of logp components~~ Implement more consistent tracking of logp components via LogJacobianAccumulator Jul 21, 2025

penelopeysm added 3 commits July 22, 2025 01:05

Fix final tests

5a4b01b

Fix docs

974c282

Fix docs/doctests

60b6863

penelopeysm added 2 commits July 22, 2025 01:18

Fix maths in LogJacobianAccumulator docstring

53a2f61

Twiddle with a comment

a47641c

penelopeysm commented Jul 22, 2025

View reviewed changes

src/transforming.jl Show resolved Hide resolved

penelopeysm commented Jul 22, 2025

View reviewed changes

penelopeysm marked this pull request as ready for review July 22, 2025 00:33

penelopeysm requested a review from mhauru July 22, 2025 00:33

Add changelog

10de51f

mhauru reviewed Jul 23, 2025

View reviewed changes

penelopeysm added 3 commits July 24, 2025 12:39

Fix accumulator docstring

a19bf2f

logJ -> logjac

48a6048

Fix logjac accumulation for StaticTransformation

049b3d3

penelopeysm requested a review from mhauru July 24, 2025 13:33

mhauru approved these changes Jul 24, 2025

View reviewed changes

mhauru merged commit bd99d4f into breaking Jul 24, 2025
20 of 21 checks passed

mhauru deleted the py/logjac branch July 24, 2025 14:05

This was referenced Jul 25, 2025

Use accumulators to fix all logp calculations when sampling TuringLang/Turing.jl#2630

Open

Release v0.37 #901

Draft

penelopeysm mentioned this pull request Jul 26, 2025

Support DPPL 0.37 TuringLang/Turing.jl#2550

Draft

		istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, vn::VarName) = istrans(vi.varinfo, vn)
		istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}) = istrans(vi.varinfo)

Implement more consistent tracking of logp components via LogJacobianAccumulator #998

Implement more consistent tracking of logp components via LogJacobianAccumulator #998

Uh oh!

Conversation

penelopeysm commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Convention: Link transform

Uh oh!

github-actions bot commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark Report for Commit 049b3d3

Computer Information

Benchmark Results

Uh oh!

codecov bot commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Jul 22, 2025

Uh oh!

penelopeysm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Implement more consistent tracking of logp components via `LogJacobianAccumulator` #998

Implement more consistent tracking of logp components via `LogJacobianAccumulator` #998

penelopeysm commented Jul 21, 2025 •

edited

Loading

github-actions bot commented Jul 21, 2025 •

edited

Loading

Benchmark Report for Commit `049b3d3`

codecov bot commented Jul 21, 2025 •

edited

Loading

penelopeysm Jul 22, 2025 •

edited

Loading