Skip to content

Accumulators, stage 1 #885

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 67 commits into from
May 2, 2025
Merged

Accumulators, stage 1 #885

merged 67 commits into from
May 2, 2025

Conversation

mhauru
Copy link
Member

@mhauru mhauru commented Apr 10, 2025

This is starting to take shape. It's too early for a review: Everything is undocumented, uncleaned, and some things are still broken. The base design is there though, and most tests pass (pointwiseloglikelihood and doctests being the exceptions), so @penelopeysm, @torfjelde, if you want to have an early look at where this is going, feel free. The most interesting files are accumulators.jl, abstract_varinfo.jl, and context_implementations.jl.

In addition to obvious things that still need doing (documentation, clean-up, new tests, adding deprecations, fixing pointwiseloglikehood), a few things I have on my mind:

  • Need to decide whether to keep the LogLikelihood and LogPrior accumulators immutable like they are now.
  • Whether getacc and similar functions should take the type of the accumulator as the index, or rather the symbol returned by accumulator_name. Leaning towards latter, but the former is what's currently implemented.
  • Maybe rename DefaultContext to AccumulationContext. Or something else? I'm not fixated on the term "accumulator".
  • Since the signature of (tilde_)assume and (tilde_)observe has changed (they no longer return logp), the whole stack of calls within tilde_obssume!! should be revisited. In particular, I'm thinking of splitting anything sampling-related to a call of tilde_obbsume with SamplingContext, that then at the end calls tilde_obssume with DefaultContext. This might be a separate PR though.
  • Benchmark
  • There are a few places where we are now unnecessarily accumulating all of log prior, log likelihood, and num produce. I should clean these up to benefit from being able to do one but not the others.
  • Make metadata.order be an accumulator as well. Probably needs to actually be in the same accumulator with NumProduce, since the two go together. Probably a separate PR though.

Edit (@penelopeysm)

Closes #580
Closes #390

penelopeysm and others added 9 commits March 5, 2025 10:34
* AbstractPPL 0.11; change prefixing behaviour

* Use DynamicPPL.prefix rather than overloading
* Unify {Untyped,Typed}{Vector,}VarInfo constructors

* Update invocations

* NTVarInfo

* Fix tests

* More fixes

* Fixes

* Fixes

* Fixes

* Use lowercase functions, don't deprecate VarInfo

* Rewrite VarInfo docstring

* Fix methods

* Fix methods (really)
Copy link
Contributor

github-actions bot commented Apr 10, 2025

Benchmark Report for Commit efc7c53

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.2 |                 1.7 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                733.9 |                34.1 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                410.5 |                45.9 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1235.0 |                27.8 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               7701.8 |                21.3 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1508.6 |                26.8 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                996.8 |                 5.0 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               6310.5 |                 3.6 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1097.4 |                 8.0 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              65578.6 |                 3.5 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8966.3 |                 9.3 |
|               Dynamic |        10 |    mooncake |             typed |   true |                136.6 |                14.0 |
|              Submodel |         1 |    mooncake |             typed |   true |                 12.0 |                 6.9 |
|                   LDA |        12 | reversediff |             typed |   true |               1182.2 |                 1.7 |

Copy link
Member

@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not reviewing actual code, just one high-level thought that struck me.

Comment on lines 121 to 125
function setlogp!!(vi::AbstractVarInfo, logp)
vi = setlogprior!!(vi, zero(logp))
vi = setloglikelihood!!(vi, logp)
return vi
end
Copy link
Member

@penelopeysm penelopeysm Apr 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking about this the other day and thought I may as well post now. The ...logp() family of functions are no longer well-defined in a world where everything is cleanly split into prior and likelihood. (only getlogp and resetlogp still make sense) I think last time we chatted about it the decision was to maybe forward the others to the likelihood methods, but I was wondering if it's actually safer to remove them (or make them error informatively) and force people to use likelihood or prior as otherwise it risks introducing subtle bugs. Backward compatibility is important but if it comes at the cost of correctness I feel kinda uneasy.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My hope was that we could deprecate them but provide the same functionality through the new functions, like above. It's a good question as to whether there are edge cases where they do not provide the same functionality. I think this is helped by the fact that PriorContext and LikelihoodContext won't exist, and hence one can't be running code where the expectation would be that ...logp() would be referring to logprior or loglikelihood in particular. And I think as long as one expects to get the logjoint out of ...logp() we can do things like above, shoving things into likelihood, and get the same results. Do you think that solves it and let's us use deprecations rather than straight-up removals, or do you see other edge cases?

Copy link
Member

@penelopeysm penelopeysm Apr 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Something like this is a case where setlogp is ill-defined:

lp = getlogp(vi_typed_metadata)
varinfos = map((
vi_untyped_metadata,
vi_untyped_vnv,
vi_typed_metadata,
vi_typed_vnv,
svi_typed,
svi_untyped,
svi_vnv,
svi_typed_ref,
svi_untyped_ref,
svi_vnv_ref,
)) do vi
# Set them all to the same values.
DynamicPPL.setlogp!!(update_values!!(vi, example_values, varnames), lp)
end

The logp here contains terms from both prior and likelihood, but after calling setlogp the prior would always be 0, which is inconsistent with the varinfo.

Of course, we can fix this on our end - we would get and set logprior and loglikelihood manually, and we can grep the codebase to make sure that there are no other ill-defined calls to setlogp. We can't guarantee that other people will be similarly careful, though (and us or anyone being careful also doesn't guarantee that everything will be fixed correctly).

Copy link
Member

@penelopeysm penelopeysm Apr 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While looking for other uses of setlogp, I encountered this:

https://github.com/TuringLang/Turing.jl/blob/fc32e10bc17ae3fda4d7e825b6fde45dc7bdb179/src/mcmc/hmc.jl#L201-L234

AdvancedHMC.Transition only contains a single notion of log density, so it's not obvious to me how we're going to extract the prior and likelihood components from it 😓 This might require upstream changes to AdvancedHMC. Since the contexts will be removed, I suspect LogDensityFunction also needs to be changed so that there's a way for it to return only the prior or the likelihood (or maybe it should return both).

(For the record, I'd be quite happy with making all of these changes!)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The logp here contains terms from both prior and likelihood, but after calling setlogp the prior would always be 0, which is inconsistent with the varinfo.

It is inconsistent, but as long as the user only uses getlogp, they would never see the difference, right? If some of logprior is accidentally stored in loglikelihood or vice versa, as long as one is using getlogp and DefaultContext that should be undetectable. What would be trouble is if someone mixes using e.g. setlogp!! and getlogprior, which would require adding calls to getlogprior after upgrading to a version that has deprecated setlogp!!, but probably people would end up doing that. Maybe the deprecation warning could say something about this?

Since the contexts will be removed, I suspect LogDensityFunction also needs to be changed so that there's a way for it to return only the prior or the likelihood (or maybe it should return both).

Yeah, this sort of stuff will come up (and is coming up) in multiple places. Anything that explicitly uses PriorContext or LikelihoodContext would need to be changed to use LogPrior and LogLikelihood accumulators instead. I'm currently doing this for pointwiselogdensities.

@mhauru
Copy link
Member Author

mhauru commented Apr 11, 2025

pointwise_logdensities now works and uses its own accumulator type rather than a context. This leaves only a handful of tests failing, for quite trivial reasons. Plenty of clean-up to do though: In fixing pointwise_logprobability I had to make substantial changes to the tilde_observe pipeline, because accumulate_observe needed to get the varname as an argument and thus had to be moved higher in the call stack. I'll have to see how to best reorganise tilde_observe in such a way that making ParticleGibbs work with it wouldn't be horrible.

Comment on lines +129 to +133
y = getindex_internal(vi, vn)
f = from_maybe_linked_internal_transform(vi, vn, dist)
x, logjac = with_logabsdet_jacobian(f, y)
vi = accumulate_assume!!(vi, x, logjac, vn, dist)
return x, vi
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How do we deal with tempering of logpdf and such now that it happens in the leaf of the call stack?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the past, we would do this by altering the logpdf coming a the assume higher up in the call tree

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are the needs of tempering? What does it need to alter?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Granted, I've only spent about 30 minutes reading about it, but I don't see the need for tempering to have such fine-grained control over the logp emitted from each individual assume / observe call -- it seems sufficient to either globally modify the logprior or loglikelihood, which can still be done with the current approach (by evaluating the model and then extracting the logp components). I had a look through the Pigeons.jl codebase for example and it doesn't seem to need to hook into the tilde pipeline. I don't want to speak for MCMCTempering since Tor is here 😉 Happy to be corrected if I'm wrong though

@penelopeysm
Copy link
Member

As for renaming DefaultContext, we chatted about it yesterday and yes I wouldn't mind calling it AccumulatorContext if there was some other context that didn't run the accumulators. (Although, I think the same effect could be achieved by just using an empty tuple of accumulators.) I think we decided to leave it to a later PR though.

Comment on lines +129 to +133
y = getindex_internal(vi, vn)
f = from_maybe_linked_internal_transform(vi, vn, dist)
x, logjac = with_logabsdet_jacobian(f, y)
vi = accumulate_assume!!(vi, x, logjac, vn, dist)
return x, vi
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Granted, I've only spent about 30 minutes reading about it, but I don't see the need for tempering to have such fine-grained control over the logp emitted from each individual assume / observe call -- it seems sufficient to either globally modify the logprior or loglikelihood, which can still be done with the current approach (by evaluating the model and then extracting the logp components). I had a look through the Pigeons.jl codebase for example and it doesn't seem to need to hook into the tilde pipeline. I don't want to speak for MCMCTempering since Tor is here 😉 Happy to be corrected if I'm wrong though

Comment on lines +441 to +443
# TODO(mhauru) How could we do this more cleanly? The problem case is map_accumulator!!
# for ThreadSafeVarInfo. In that one, if the map produces e.g a ForwardDiff.Dual, but
# the accumulators in the VarInfo are plain floats, we error since we can't change the
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we're almost there, just a few things more to look at.

I have stared at this for quite a while and I understand the issue with map_accumulator!!, but I don't get how this relates to unflatten?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comes about in situations where you've made a regular VarInfo where the log probs are floats, and then you e.g. give it to a LogDensityFunction which wants to evaluate your model with values that are of some float-look-a-like type such as ForwardDiff.Dual. This will manifest itself by unflatten getting called with a vector of Duals. Further down the line this will cause us to add stuff to log probs that is also Duals.

Ideally I would like a situation where at the point when I have e.g. a LogPrior{Float64} and I want to add to it a LogPrior{Dual{Float64}} this will naturally result in a LogPrior{Dual{Float64}} (because we defer to the underlying type when adding LogPriors, and they know how to handle Float64 + Dual{Float64}) and then that gets assigned to vi.accs, making use of the fact that since all of our functions for accs are !! we can make the assignment not-in-place if need be. How elegant!

But ThreadSafeVarInfo rains on this parade by being unable to change the types of the accumulators, because it needs to update each sub-AccumulatorTuple for each thread independently. Grrrr.

The solution is to do what we used to do, which is to say already in unflatten, before any splitting into multiple threads happens, that we force our log probs to be of the same element type as the vector x that was given to unflatten. This is crude, and for instance means that if you deliberately made your x be, say, Float32s and your log probs be Float64s, that's ruined here. The only other solution I could think of was to special case in unflatten to only do this conversion if the element type is one of the AD tracer types, but then that would require enumerating them all and having all the AD packages as dependencies.

(I think the real solution is to make ThreadSafeVarInfo better, maybe with locks, but that's a different story, and may not be worth the effort.)

Does that help? I'll improve the comment once I know how to best improve it.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Several thoughts:

  1. If the problem behaviour is restricted to TSVI, could we keep the special-casing to TSVI e.g. with a special method for unflatten?
  2. In the long run it seems that we need to have some recognition that T = ForwardDiff.Dual{Float64} (i.e. evaluating grad(logp)) is orthogonal to from T = Float32 (i.e. evaluating logp with different precision) -- and they could even be mix-and-matched i.e. T = ForwardDiff.Dual{Float32} which means evaluating grad(logp) with different precision. convert_eltype isn't right now clever enough to make this differentiation, but it appears to me that we indeed do need to (and should) hardcode this. The lowest effort way seems to be to use extensions. I wonder if in the long term this is something that could go into DifferentiationInterface.
  3. If it can't be solved easily, could you open an issue to track this?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the problem behaviour is restricted to TSVI, could we keep the special-casing to TSVI e.g. with a special method for unflatten?

I suspect we might have situations where unflatten is called first on a regular VI and only then it's converted to TSVI.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See #906 for more on this.

Copy link
Member

@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've looked through the rest of the files and, okay, I'll admit I haven't minutely checked every single line, but I did read and I don't think I saw anything that merited closer inspection.

Final comment is about the performance. Have you made any inroads into finding out why the small models are a lot worse? Does it matter? Should the mutable-accumulator investigation be a separate PR?

@mhauru
Copy link
Member Author

mhauru commented May 1, 2025

I figured out the slowness issue: unflatten was type unstable for silly reasons. With that fixed, here are benchmarks on just the simplest benchmark model with various VarInfo types. On this PR:

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                  4.4 |                 1.8 |
| Simple assume observe |         1 | forwarddiff |           untyped |  false |                 27.4 |                 2.0 |
| Simple assume observe |         1 | forwarddiff | simple_namedtuple |  false |                  0.7 |                 6.0 |
| Simple assume observe |         1 | forwarddiff |       simple_dict |  false |                  4.4 |                 1.8 |

On v0.36.1:

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                  7.9 |                 1.8 |
| Simple assume observe |         1 | forwarddiff |           untyped |  false |                 31.6 |                 3.4 |
| Simple assume observe |         1 | forwarddiff | simple_namedtuple |  false |                  1.6 |                 3.0 |
| Simple assume observe |         1 | forwarddiff |       simple_dict |  false |                  6.3 |                 2.0 |

The type instability was there already before, its consequences just got more severe with accumulators. I added a test that would have caught it.

The whole usual benchmark suite, on this PR:

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                 10.8 |                 1.7 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                543.4 |                43.1 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                235.5 |                79.3 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |                899.0 |                34.2 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               5672.8 |                22.9 |
|           Smorgasbord |       201 | reversediff |             typed |   true |                920.6 |                23.5 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                948.3 |                 4.2 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5229.5 |                 4.6 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |                749.7 |                 7.9 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              57145.2 |                 4.5 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               6160.8 |                 9.0 |
|               Dynamic |        10 |    mooncake |             typed |   true |                152.4 |                 7.3 |
|              Submodel |         1 |    mooncake |             typed |   true |                 16.9 |                 7.6 |
|                   LDA |        12 | reversediff |             typed |   true |                917.5 |                 1.4 |

For contrast, copypasta of the results on current main from above:

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                 12.6 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                737.1 |                36.2 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                244.6 |                74.2 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1161.7 |                30.6 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               3284.7 |                22.4 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1108.1 |                22.0 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1157.0 |                 3.6 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5365.1 |                 5.0 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |                749.8 |                 7.9 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              56613.4 |                 5.0 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               6427.4 |                 8.8 |
|               Dynamic |        10 |    mooncake |             typed |   true |                145.2 |                 7.7 |
|              Submodel |         1 |    mooncake |             typed |   true |                 17.4 |                 7.8 |
|                   LDA |        12 | reversediff |             typed |   true |                460.9 |                 2.8 |

Mildly curious about what's going on with LDA (it's a horrible model anyway), and I also still want to try the mutable accumulators thing, but mostly I think we are done with performance concerns.

@mhauru
Copy link
Member Author

mhauru commented May 1, 2025

Making accumulators mutable seemed to harm speed, as did most of the @inline statements I tried.

@mhauru
Copy link
Member Author

mhauru commented May 1, 2025

For the renaming of DefaultContext, both @penelopeysm and I thought it's better to do it later if we do it, once the more general restructuring of the tilde pipeline is done.

Copy link
Member

@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My only last thought was that user-facing docs would be v useful, although I'm guessing it's already on your list.

Comment on lines +441 to +443
# TODO(mhauru) How could we do this more cleanly? The problem case is map_accumulator!!
# for ThreadSafeVarInfo. In that one, if the map produces e.g a ForwardDiff.Dual, but
# the accumulators in the VarInfo are plain floats, we error since we can't change the
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Several thoughts:

  1. If the problem behaviour is restricted to TSVI, could we keep the special-casing to TSVI e.g. with a special method for unflatten?
  2. In the long run it seems that we need to have some recognition that T = ForwardDiff.Dual{Float64} (i.e. evaluating grad(logp)) is orthogonal to from T = Float32 (i.e. evaluating logp with different precision) -- and they could even be mix-and-matched i.e. T = ForwardDiff.Dual{Float32} which means evaluating grad(logp) with different precision. convert_eltype isn't right now clever enough to make this differentiation, but it appears to me that we indeed do need to (and should) hardcode this. The lowest effort way seems to be to use extensions. I wonder if in the long term this is something that could go into DifferentiationInterface.
  3. If it can't be solved easily, could you open an issue to track this?

@mhauru mhauru merged commit 299e17b into breaking May 2, 2025
17 of 18 checks passed
@mhauru mhauru deleted the mhauru/custom-accumulators branch May 2, 2025 13:12
@mhauru
Copy link
Member Author

mhauru commented May 2, 2025

Thanks for the great feedback!

github-merge-queue bot pushed a commit that referenced this pull request Aug 7, 2025
* Bump minor version to 0.37.0

* Accumulators, stage 1 (#885)

* Release 0.36

* AbstractPPL 0.11 + change prefixing behaviour (#830)

* AbstractPPL 0.11; change prefixing behaviour

* Use DynamicPPL.prefix rather than overloading

* Remove VarInfo(VarInfo, params) (#870)

* Unify `{untyped,typed}_{vector_,}varinfo` constructor functions (#879)

* Unify {Untyped,Typed}{Vector,}VarInfo constructors

* Update invocations

* NTVarInfo

* Fix tests

* More fixes

* Fixes

* Fixes

* Fixes

* Use lowercase functions, don't deprecate VarInfo

* Rewrite VarInfo docstring

* Fix methods

* Fix methods (really)

* Draft of accumulators

* Fix some variable names

* Fix pointwise_logdensities, gut tilde_observe, remove resetlogp!!

* Map rather than broadcast

Co-authored-by: Tor Erlend Fjelde <[email protected]>

* Start documenting accumulators

* Use Val{symbols} instead of AccTypes to index

* More documentation for accumulators

* Link varinfo by default in AD testing utilities; make test suite run on linked varinfos (#890)

* Link VarInfo by default

* Tweak interface

* Fix tests

* Fix interface so that callers can inspect results

* Document

* Fix tests

* Fix changelog

* Test linked varinfos

Closes #891

* Fix docstring + use AbstractFloat

* Fix resetlogp!! and type stability for accumulators

* Fix type rigidity of LogProbs and NumProduce

* Fix uses of getlogp and other assorted issues

* setaccs!! nicer interface and logdensity function fixes

* Revert back to calling the macro @addlogprob!

* Remove a dead test

* Clarify a comment

* Implement split/combine for PointwiseLogdensityAccumulator

* Switch ThreadSafeVarInfo.accs_by_thread to be a tuple

* Fix `condition` and `fix` in submodels (#892)

* Fix conditioning in submodels

* Simplify contextual_isassumption

* Add documentation

* Fix some tests

* Add tests; fix a bunch of nested submodel issues

* Fix fix as well

* Fix doctests

* Add unit tests for new functions

* Add changelog entry

* Update changelog

Co-authored-by: Hong Ge <[email protected]>

* Finish docs

* Add a test for conditioning submodel via arguments

* Clean new tests up a bit

* Fix for VarNames with non-identity lenses

* Apply suggestions from code review

Co-authored-by: Markus Hauru <[email protected]>

* Apply suggestions from code review

* Make PrefixContext contain a varname rather than symbol (#896)

---------

Co-authored-by: Hong Ge <[email protected]>
Co-authored-by: Markus Hauru <[email protected]>

* Revert ThreadSafeVarInfo back to Vectors and fix some AD type casting in (Simple)VarInfo

* Improve accumulator docs

* Add test/accumulators.jl

* Docs fixes

* Various small fixes

* Make DynamicTransformation not use accumulators other than LogPrior

* Fix variable order and name of map_accumulator!!

* Typo fixing

* Small improvement to ThreadSafeVarInfo

* Fix demo_dot_assume_observe_submodel prefixing

* Typo fixing

* Miscellaneous small fixes

* HISTORY entry and more miscellanea

* Add more tests for accumulators

* Improve accumulators docstrings

* Fix a typo

* Expand HISTORY entry

* Add accumulators to API docs

* Remove unexported functions from API docs

* Add NamedTuple methods for get/set/acclogp

* Fix setlogp!! with single scalar to error

* Export AbstractAccumulator, fix a docs typo

* Apply suggestions from code review

Co-authored-by: Penelope Yong <[email protected]>

* Rename LogPrior -> LogPriorAccumulator, and Likelihood and NumProduce

* Type bound log prob accumulators with T<:Real

* Add @addlogprior! and @addloglikelihood!

* Apply suggestions from code review

Co-authored-by: Penelope Yong <[email protected]>

* Move default accumulators to default_accumulators.jl

* Fix some tests

* Introduce default_accumulators()

* Go back to only having @addlogprob!

* Fix tilde_observe!! prefixing

* Fix default_accumulators internal type

* Make unflatten more type stable, and add a test for it

* Always print all benchmark results

* Move NumProduce VI functions to abstract_varinfo.jl

---------

Co-authored-by: Penelope Yong <[email protected]>
Co-authored-by: Tor Erlend Fjelde <[email protected]>
Co-authored-by: Hong Ge <[email protected]>

* Replace PriorExtractorContext with PriorDistributionAccumulator (#907)

* Implement values_as_in_model using an accumulator (#908)

* Implement values_as_in_model using an accumulator

* Make make_varname_expression a function

* Refuse to combine ValuesAsInModelAccumulators with different include_colon_eqs

* Fix nested context test

* Bump DynamicPPL versions

* Fix merge (1)

* Add benchmark Pkg source

* [no ci] Don't need to dev again

* Disable use_closure for ReverseDiff

* Revert "Disable use_closure for ReverseDiff"

This reverts commit 3cb47cd.

* Fix LogDensityAt struct

* Try not duplicating

* Update comment pointing to closure benchmarks

* Remove `context` from model evaluation (use `model.context` instead) (#952)

* Change `evaluate!!` API, add `sample!!`

* Fix literally everything else that I broke

* Fix some docstrings

* fix ForwardDiffExt (look, multiple dispatch bad...)

* Changelog

* fix a test

* Fix docstrings

* use `sample!!`

* Fix a couple more cases

* Globally rename `sample!!` -> `evaluate_and_sample!!`, add changelog warning

* Mark function as Const for Enzyme tests (#957)

* Move submodel code to submodel.jl; remove `@submodel` (#959)

* Move submodel code to submodel.jl

* Remove `@submodel`

* Fix missing field tests for 1.12 (#961)

* Remove 3-argument `{_,}evaluate!!`; clean up submodel code (#960)

* Clean up submodel code, remove 3-arg `_evaluate!!`

* Remove 3-argument `evaluate!!` as well

* Update changelog

* Improve submodel error message

* Fix doctest

* Add error hint for three-argument evaluate!!

* Improve API for AD testing (#964)

* Rework API for AD testing

* Fix test

* Add `rng` keyword argument

* Use atol and rtol

* remove unbound type parameter (?)

* Don't need to do elementwise check

* Update changelog

* Fix typo

* DebugAccumulator (plus tiny bits and pieces) (#976)

* DebugContext -> DebugAccumulator

* Changelog

* Force `conditioned` to return a dict

* fix conditioned implementation

* revert `conditioned` bugfix (will merge this to main instead)

* fix show

* Fix doctests

* fix doctests 2

* Make VarInfo actually mandatory in check_model

* Re-implement `missing` check

* Revert `combine` signature in docstring

* Revert changes to `Base.show` on AccumulatorTuple

* Add TODO comment about VariableOrderAccumulator

Co-authored-by: Markus Hauru <[email protected]>

* Fix doctests

---------

Co-authored-by: Markus Hauru <[email protected]>

* VariableOrderAccumulator (#940)

* Turn NumProduceAccumulator into VariableOrderAccumulator

* Add comparison methods

* Make VariableOrderAccumulator use regular Dict

* Use copy rather than deepcopy for accumulators

* Minor docstring touchup

* Remove unnecessary use of NumProduceAccumulator

* Fix split(VariableOrderAccumulator)

* Remove NumProduceAcc from Debug

* Fix set_retained_vns_del!

---------

Co-authored-by: Penelope Yong <[email protected]>

* Accumulators stage 2 (#925)

* Give LogDensityFunction the getlogdensity field

* Allow missing LogPriorAccumulator when linking

* Trim whitespace

* Run formatter

* Fix a few typos

* Fix comma -> semicolon

* Fix `LogDensityAt` invocation

* Fix one last test

* Fix tests

---------

Co-authored-by: Penelope Yong <[email protected]>

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

* logjac accumulator

* Fix tests

* Fix a whole bunch of stuff

* Fix final tests

* Fix docs

* Fix docs/doctests

* Fix maths in LogJacobianAccumulator docstring

* Twiddle with a comment

* Add changelog

* Fix accumulator docstring

* logJ -> logjac

* Fix logjac accumulation for StaticTransformation

* Fix behaviour of `set_retained_vns_del!` for `num_produce == 0` (#1000)

* `InitContext`, part 2 - Move `hasvalue` and `getvalue` to AbstractPPL; enforce key type of `AbstractDict` (#980)

* point to unmerged AbstractPPL branch

* Remove code that was moved to AbstractPPL

* Remove Dictionaries with Any key type

* Fix bad merge conflict resolution

* Fix doctests

* Point to [email protected]

This reverts commit 709dc9e.

* Fix doctests

* Fix docs AbstractPPL bound

* Remove stray `Pkg.update()`

* Accumulator miscellanea: Subset, merge, acclogp, and LogProbAccumulator (#999)

* logjac accumulator

* Fix tests

* Fix a whole bunch of stuff

* Fix final tests

* Fix docs

* Fix docs/doctests

* Fix maths in LogJacobianAccumulator docstring

* Twiddle with a comment

* Add changelog

* Simplify accs with LogProbAccumulator

* Replace + with accumulate for LogProbAccs

* Introduce merge and subset for accs

* Improve acc tests

* Fix docstring typo.

Co-authored-by: Penelope Yong <[email protected]>

* Fix merge

---------

Co-authored-by: Penelope Yong <[email protected]>

* Minor tweak to changelog wording

---------

Co-authored-by: Penelope Yong <[email protected]>
Co-authored-by: Tor Erlend Fjelde <[email protected]>
Co-authored-by: Hong Ge <[email protected]>
@penelopeysm penelopeysm mentioned this pull request Aug 7, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants