Accumulators stage 2

[mhauru]

I'm making various changes to accumulators as I integrate Turing.jl with them. This isn't ready for review yet, but @penelopeysm, can I ask you for your thoughts on the design of the proposed changes to LogDensityFunction? It would now have a new field, given as the second constructor argument, called getlogdensity. By default it's getlogjoint, but you can set it to getlogprior or getloglikelihood or any other function that takes an AbstractVarInfo. Its return value will be the return value of logdensity_and_gradient etc. The default VarInfo, if one isn't given by the user, is also now set based on getlogdensity, to make sure it has the right accumulators.

[github-actions]

Benchmark Report for Commit 175b633

Computer Information

Julia Version 1.11.5
Commit 760b2e5b739 (2025-04-14 06:53 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 |                  8.6 |                 1.8 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                690.4 |                35.9 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                450.5 |                45.4 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1271.2 |                27.3 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               8242.7 |                19.6 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1546.8 |                26.0 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1012.4 |                 4.9 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5831.2 |                 3.9 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1039.8 |                 8.6 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              64818.4 |                 3.7 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8737.8 |                 9.8 |
|               Dynamic |        10 |    mooncake |             typed |   true |                145.0 |                11.8 |
|              Submodel |         1 |    mooncake |             typed |   true |                 13.1 |                 6.4 |
|                   LDA |        12 | reversediff |             typed |   true |               1199.1 |                 1.8 |

[codecov]

Codecov Report

Attention: Patch coverage is 65.51724% with 10 lines in your changes missing coverage. Please review.

Please upload report for BASE (breaking@d4ef1f2). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/logdensityfunction.jl	69.23%	4 Missing ⚠️
src/simple_varinfo.jl	0.00%	4 Missing ⚠️
src/test_utils/ad.jl	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             breaking     #925   +/-   ##
===========================================
  Coverage            ?   81.11%           
===========================================
  Files               ?       37           
  Lines               ?     4035           
  Branches            ?        0           
===========================================
  Hits                ?     3273           
  Misses              ?      762           
  Partials            ?        0           

☔ 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.

[penelopeysm]

Hmmm, in terms of design, the other potential option I see is to make the new getlogdensity a keyword argument (which means that varinfo can't depend on it)

user specifies...	getlogdensity as positional argument	getlogdensity as keyword argument
+getlogdensity +varinfo	have to make sure they are consistent	have to make sure they are consistent
+getlogdensity -varinfo	generates the most efficient varinfo	uses default varinfo, potentially inefficient
-getlogdensity +varinfo	not possible	preserves current behaviour
-getlogdensity -varinfo	preserves current behaviour	preserves current behaviour

So it seems the main tradeoff is that having it as a positional argument allows for the most efficient varinfo (one without a likelihood accumulator) to be generated if you do e.g. LogDensityFunction(model, getlogprior), but in return, you can't call LogDensityFunction(model, varinfo), you have to do LogDensityFunction(model, getlogjoint, varinfo).

I think my first instinct is that I prefer the latter, and you could either insert a check in the inner constructor to see if the accumulators are consistent (i.e. warn if you have extra accs, or error if you don't have enough accs), or just straight-up call setaccs in the inner constructor to force consistency. I'm generally not a huge fan of Julia's optional positional arguments because they aren't truly optional – you can only omit them starting from the end (unless you start declaring 2^N methods).

Do you prefer the former? I'll mull over it a bit more too and see if any ideas come up.

[mhauru]

I wouldn't even bother checking for consistency if the user provides both getlogdensity and varinfo. If you're getting that hands-on with your LogDensityFunction I assume you know what you're doing and if not that's on you. If there's a mismatch, you'll probably just get an error saying "no accumulator X available in your VarInfo".

I would be happy to give up the case of LogDensityFunction(model, varinfo) and force people to do LogDensityFunction(model, getlogjoint, varinfo) because giving the varinfo argument I think means you're dealing with some pretty low-level stuff, most likely either trying to use a non-standard subtype of AbstractVarInfo or a linked VarInfo. If you're thinking about that sorts of things, I think having to give it a few seconds and one more line of code to specify which log probability you are after (a much more "user-level"/statistical concern) seems fine to me. (Similarly this is why I think the context should be the last argument, as it is now, because that's getting even deeper into the internals of DPPL.)

There's also something about type dispatch treating keyword args differently. Sometimes methods are specialised on them and sometimes they are not, and I don't understand the details.