Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` #1129

penelopeysm · 2025-11-08T17:50:48Z

This has no real impact on DynamicPPL, but using this in Turing should lead to some nice speedups because this:

is faster to evaluate than regular unflatten + evaluate!!;
avoids deepcopying VarInfo.

Closes #1119.

Since the intention of this is to replace Turing.Inference.Transition with ParamsWithStats, I also added the necessary AbstractMCMC.bundle_samples method here. MCMCChainsExt is the natural place for it to live (it could be defined in Turing, but that would be piracy).

github-actions · 2025-11-08T17:55:30Z

Benchmark Report

this PR's head: 0a9557ec0633bb85a42b0124727b7bfe8de82e12
base branch: 3cd8d3431e14ebc581266c1323d1db8a5bd4c0eb

Computer Information

Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 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

┌───────────────────────┬───────┬─────────────┬───────────────────┬────────┬─────────────────────────────────┬───────────────────────────┬──────────────────────────────────┐
│                       │       │             │                   │        │        t(eval) / t(ref)         │     t(grad) / t(eval)     │         t(grad) / t(ref)         │
│                       │       │             │                   │        │ ──────────┬───────────┬──────── │ ──────┬─────────┬──────── │ ──────────┬────────────┬──────── │
│                 Model │   Dim │  AD Backend │           VarInfo │ Linked │      base │   this PR │ speedup │  base │ this PR │ speedup │      base │    this PR │ speedup │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│               Dynamic │    10 │    mooncake │             typed │   true │    504.53 │    519.32 │    0.97 │  7.67 │    7.35 │    1.04 │   3870.25 │    3819.24 │    1.01 │
│                   LDA │    12 │ reversediff │             typed │   true │   2915.66 │   2983.56 │    0.98 │  2.27 │    2.12 │    1.07 │   6611.06 │    6319.26 │    1.05 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │ 162397.29 │ 163322.69 │    0.99 │  5.90 │    6.35 │    0.93 │ 957822.51 │ 1036552.28 │    0.92 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │  14807.28 │  14817.98 │    1.00 │  6.12 │    6.01 │    1.02 │  90690.53 │   89119.79 │    1.02 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │  45304.69 │  47745.23 │    0.95 │  7.62 │    7.47 │    1.02 │ 345166.70 │  356839.74 │    0.97 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │   5185.70 │   9561.27 │    0.54 │  6.77 │    4.17 │    1.62 │  35121.79 │   39884.18 │    0.88 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │     20.42 │     18.86 │    1.08 │  1.72 │    1.91 │    0.90 │     35.08 │      35.94 │    0.98 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │   2990.16 │   3078.52 │    0.97 │ 47.17 │   44.42 │    1.06 │ 141046.21 │  136747.70 │    1.03 │
│           Smorgasbord │   201 │ forwarddiff │       simple_dict │   true │  25176.26 │  26895.89 │    0.94 │ 28.33 │   28.60 │    0.99 │ 713344.90 │  769187.00 │    0.93 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │ simple_namedtuple │   true │   1305.72 │   1308.92 │    1.00 │ 81.47 │   81.82 │    1.00 │ 106373.25 │  107092.84 │    0.99 │
│           Smorgasbord │   201 │      enzyme │             typed │   true │   2943.02 │   3020.11 │    0.97 │  4.70 │    4.76 │    0.99 │  13840.55 │   14372.07 │    0.96 │
│           Smorgasbord │   201 │    mooncake │             typed │   true │   2935.20 │   3052.52 │    0.96 │  4.93 │    4.92 │    1.00 │  14464.31 │   15021.72 │    0.96 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│           Smorgasbord │   201 │ reversediff │             typed │   true │   2959.98 │   3091.69 │    0.96 │ 57.09 │   55.54 │    1.03 │ 168991.72 │  171717.17 │    0.98 │
│           Smorgasbord │   201 │ forwarddiff │      typed_vector │   true │   3033.14 │   3099.62 │    0.98 │ 77.26 │   48.93 │    1.58 │ 234328.22 │  151663.75 │    1.55 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │   2600.79 │   2620.72 │    0.99 │ 51.58 │   49.99 │    1.03 │ 134142.57 │  131022.47 │    1.02 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼───────┼─────────┼─────────┼───────────┼────────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │    untyped_vector │   true │   2696.82 │   2721.07 │    0.99 │ 49.60 │   51.44 │    0.96 │ 133755.53 │  139972.91 │    0.96 │
│              Submodel │     1 │    mooncake │             typed │   true │     28.12 │     28.85 │    0.97 │  5.14 │    6.64 │    0.78 │    144.64 │     191.46 │    0.76 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴───────────┴───────────┴─────────┴───────┴─────────┴─────────┴───────────┴────────────┴─────────┘

github-actions · 2025-11-08T18:01:59Z

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

codecov · 2025-11-08T18:12:01Z

Codecov Report

❌ Patch coverage is 72.72727% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.59%. Comparing base (3cd8d34) to head (0a9557e).
⚠️ Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
ext/DynamicPPLMCMCChainsExt.jl	0.00%	9 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #1129      +/-   ##
============================================
- Coverage     81.67%   81.59%   -0.09%     
============================================
  Files            42       42              
  Lines          3930     3955      +25     
============================================
+ Hits           3210     3227      +17     
- Misses          720      728       +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.

penelopeysm · 2025-11-21T11:50:59Z

ping @sunxd3 :)

sunxd3

looks good, couple of tiny questions

sunxd3 · 2025-11-21T14:33:37Z

src/chains.jl

+            DynamicPPL.LogPriorAccumulator(),
+            DynamicPPL.LogLikelihoodAccumulator(),
+            DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),


I am looking at

DynamicPPL.jl/src/fasteval.jl

Lines 247 to 262 in 5d4527a

"""

fast_ldf_accs(getlogdensity::Function)

Determine which accumulators are needed for fast evaluation with the given

`getlogdensity` function.

"""

fast_ldf_accs(::Function) = default_accumulators()

fast_ldf_accs(::typeof(getlogjoint_internal)) = default_accumulators()

function fast_ldf_accs(::typeof(getlogjoint))

return AccumulatorTuple((LogPriorAccumulator(), LogLikelihoodAccumulator()))

end

function fast_ldf_accs(::typeof(getlogprior_internal))

return AccumulatorTuple((LogPriorAccumulator(), LogJacobianAccumulator()))

end

fast_ldf_accs(::typeof(getlogprior)) = AccumulatorTuple((LogPriorAccumulator(),))

fast_ldf_accs(::typeof(getloglikelihood)) = AccumulatorTuple((LogLikelihoodAccumulator(),))

and wonder if there are time we need the LogJacobianAccumulator

Ahh, right. I think it doesn't matter for the present purposes because we always want the output to not include the Jacobian term i.e. logprior and logjoint are 'as seen in the model'.

ext/DynamicPPLMCMCChainsExt.jl

* v0.39 * Update DPPL compats for benchmarks and docs * remove merge conflict markers * Remove `NodeTrait` (#1133) * Remove NodeTrait * Changelog * Fix exports * docs * fix a bug * Fix doctests * Fix test * tweak changelog * FastLDF / InitContext unified (#1132) * Fast Log Density Function * Make it work with AD * Optimise performance for identity VarNames * Mark `get_range_and_linked` as having zero derivative * Update comment * make AD testing / benchmarking use FastLDF * Fix tests * Optimise away `make_evaluate_args_and_kwargs` * const func annotation * Disable benchmarks on non-typed-Metadata-VarInfo * Fix `_evaluate!!` correctly to handle submodels * Actually fix submodel evaluate * Document thoroughly and organise code * Support more VarInfos, make it thread-safe (?) * fix bug in parsing ranges from metadata/VNV * Fix get_param_eltype for TSVI * Disable Enzyme benchmark * Don't override _evaluate!!, that breaks ForwardDiff (sometimes) * Move FastLDF to experimental for now * Fix imports, add tests, etc * More test fixes * Fix imports / tests * Remove AbstractFastEvalContext * Changelog and patch bump * Add correctness tests, fix imports * Concretise parameter vector in tests * Add zero-allocation tests * Add Chairmarks as test dep * Disable allocations tests on multi-threaded * Fast InitContext (#1125) * Make InitContext work with OnlyAccsVarInfo * Do not convert NamedTuple to Dict * remove logging * Enable InitFromPrior and InitFromUniform too * Fix `infer_nested_eltype` invocation * Refactor FastLDF to use InitContext * note init breaking change * fix logjac sign * workaround Mooncake segfault * fix changelog too * Fix get_param_eltype for context stacks * Add a test for threaded observe * Export init * Remove dead code * fix transforms for pathological distributions * Tidy up loads of things * fix typed_identity spelling * fix definition order * Improve docstrings * Remove stray comment * export get_param_eltype (unfortunatley) * Add more comment * Update comment * Remove inlines, fix OAVI docstring * Improve docstrings * Simplify InitFromParams constructor * Replace map(identity, x[:]) with [i for i in x[:]] * Simplify implementation for InitContext/OAVI * Add another model to allocation tests Co-authored-by: Markus Hauru <markus@mhauru.org> * Revert removal of dist argument (oops) * Format * Update some outdated bits of FastLDF docstring * remove underscores --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * implement `LogDensityProblems.dimension` * forgot about capabilities... * use interpolation in run_ad * Improvements to benchmark outputs (#1146) * print output * fix * reenable * add more lines to guide the eye * reorder table * print tgrad / trel as well * forgot this type * Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` (#1129) * Implement `ParamsWithStats` for `FastLDF` * Add comments * Implement `bundle_samples` for ParamsWithStats -> MCMCChains * Remove redundant comment * don't need Statistics? * Make FastLDF the default (#1139) * Make FastLDF the default * Add miscellaneous LogDensityProblems tests * Use `init!!` instead of `fast_evaluate!!` * Rename files, rebalance tests * Implement `predict`, `returned`, `logjoint`, ... with `OnlyAccsVarInfo` (#1130) * Use OnlyAccsVarInfo for many re-evaluation functions * drop `fast_` prefix * Add a changelog * Improve FastLDF type stability when all parameters are linked or unlinked (#1141) * Improve type stability when all parameters are linked or unlinked * fix a merge conflict * fix enzyme gc crash (locally at least) * Fixes from review * Make threadsafe evaluation opt-in (#1151) * Make threadsafe evaluation opt-in * Reduce number of type parameters in methods * Make `warned_warn_about_threads_threads_threads_threads` shorter * Improve `setthreadsafe` docstring * warn on bare `@threads` as well * fix merge * Fix performance issues * Use maxthreadid() in TSVI * Move convert_eltype code to threadsafe eval function * Point to new Turing docs page * Add a test for setthreadsafe * Tidy up check_model * Apply suggestions from code review Fix outdated docstrings Co-authored-by: Markus Hauru <markus@mhauru.org> * Improve warning message * Export `requires_threadsafe` * Add an actual docstring for `requires_threadsafe` --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * Standardise `:lp` -> `:logjoint` (#1161) * Standardise `:lp` -> `:logjoint` * changelog * fix a test --------- Co-authored-by: Markus Hauru <mhauru@turing.ac.uk> Co-authored-by: Markus Hauru <markus@mhauru.org>

penelopeysm changed the base branch from main to py/fastldf November 8, 2025 17:50

github-actions bot assigned penelopeysm Nov 8, 2025

penelopeysm mentioned this pull request Nov 8, 2025

Path to FastLDF #1120

Closed

penelopeysm force-pushed the py/params-from-ldf branch from c475f96 to 92cbb20 Compare November 13, 2025 16:46

penelopeysm marked this pull request as ready for review November 13, 2025 16:46

penelopeysm changed the base branch from py/fastldf to breaking November 13, 2025 16:46

penelopeysm force-pushed the py/params-from-ldf branch from 92cbb20 to 39df719 Compare November 13, 2025 16:46

penelopeysm changed the title ~~Allow generation of ParamsWithStats from FastLDF plus parameters~~ Allow generation of ParamsWithStats from FastLDF plus parameters, and also bundle_samples Nov 13, 2025

penelopeysm requested a review from sunxd3 November 13, 2025 18:46

penelopeysm mentioned this pull request Nov 13, 2025

Turing compatibility nsiccha/NUTS.jl#1

Open

penelopeysm force-pushed the py/params-from-ldf branch from 3e646ca to 0dd2f70 Compare November 14, 2025 00:35

penelopeysm added 2 commits November 14, 2025 00:40

Implement ParamsWithStats for FastLDF

fa32c88

Add comments

2bf5b18

penelopeysm force-pushed the py/params-from-ldf branch from 0dd2f70 to 077457f Compare November 14, 2025 00:40

Implement bundle_samples for ParamsWithStats -> MCMCChains

2fad97b

penelopeysm force-pushed the py/params-from-ldf branch from 0150ef4 to 2fad97b Compare November 14, 2025 02:06

penelopeysm added 2 commits November 18, 2025 01:42

Merge branch 'breaking' into py/params-from-ldf

8129340

Merge branch 'breaking' into py/params-from-ldf

5d4527a

sunxd3 reviewed Nov 21, 2025

View reviewed changes

Remove redundant comment

82616cd

penelopeysm requested a review from sunxd3 November 21, 2025 18:59

Merge branch 'breaking' into py/params-from-ldf

60dbab1

sunxd3 approved these changes Nov 21, 2025

View reviewed changes

don't need Statistics?

0a9557e

penelopeysm merged commit 4a11560 into breaking Nov 22, 2025
19 of 21 checks passed

penelopeysm deleted the py/params-from-ldf branch November 22, 2025 00:26

penelopeysm mentioned this pull request Nov 25, 2025

Implement ParamsWithStats(::FastLDF, ::AbstractVector{<:Real}) #1119

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` #1129

Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` #1129

penelopeysm commented Nov 8, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Nov 8, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Nov 8, 2025

Uh oh!

codecov bot commented Nov 8, 2025 •

edited

Loading

Uh oh!

penelopeysm commented Nov 21, 2025

Uh oh!

sunxd3 left a comment

Uh oh!

sunxd3 Nov 21, 2025

Uh oh!

penelopeysm Nov 21, 2025

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	"""
	fast_ldf_accs(getlogdensity::Function)

	Determine which accumulators are needed for fast evaluation with the given
	`getlogdensity` function.
	"""
	fast_ldf_accs(::Function) = default_accumulators()
	fast_ldf_accs(::typeof(getlogjoint_internal)) = default_accumulators()
	function fast_ldf_accs(::typeof(getlogjoint))
	return AccumulatorTuple((LogPriorAccumulator(), LogLikelihoodAccumulator()))
	end
	function fast_ldf_accs(::typeof(getlogprior_internal))
	return AccumulatorTuple((LogPriorAccumulator(), LogJacobianAccumulator()))
	end
	fast_ldf_accs(::typeof(getlogprior)) = AccumulatorTuple((LogPriorAccumulator(),))
	fast_ldf_accs(::typeof(getloglikelihood)) = AccumulatorTuple((LogLikelihoodAccumulator(),))

Allow generation of ParamsWithStats from FastLDF plus parameters, and also bundle_samples #1129

Allow generation of ParamsWithStats from FastLDF plus parameters, and also bundle_samples #1129

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Benchmark Report

Computer Information

Benchmark Results

Uh oh!

github-actions bot commented Nov 8, 2025

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

penelopeysm commented Nov 21, 2025

Uh oh!

sunxd3 left a comment

Choose a reason for hiding this comment

Uh oh!

sunxd3 Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

penelopeysm Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` #1129

Allow generation of `ParamsWithStats` from `FastLDF` plus parameters, and also `bundle_samples` #1129

penelopeysm commented Nov 8, 2025 •

edited

Loading

github-actions bot commented Nov 8, 2025 •

edited

Loading

codecov bot commented Nov 8, 2025 •

edited

Loading