TuringLang · mhauru · May 2, 2025 · Mar 5, 2025 · Mar 20, 2025 · Mar 22, 2025
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,5 +1,23 @@
 # DynamicPPL Changelog
 
+## 0.37.0
+
+**Breaking changes**
+
+### Accumulators
+
+This release overhauls how VarInfo objects track variables such as the log joint probability. The new approach is to use what we call accumulators: Objects that the VarInfo carries on it that may change their state at each `tilde_assume!!` and `tilde_observe!!` call based on the value of the variable in question. They replace both variables that were previously hard-coded in the `VarInfo` object (`logp` and `num_produce`) and some contexts. This brings with it a number of breaking changes:
+
+  - `PriorContext` and `LikelihoodContext` no longer exist. By default, a `VarInfo` tracks both the log prior and the log likelihood separately, and they can be accessed with `getlogprior` and `getloglikelihood`. If you want to execute a model while only accumulating one of the two (to save clock cycles), you can do so by creating a `VarInfo` that only has one accumulator in it, e.g. `varinfo = setaccs!!(varinfo, (LogPriorAccumulator(),))`.
+  - `MiniBatchContext` does not exist anymore. It can be replaced by creating and using a custom accumulator that replaces the default `LikelihoodContext`. We may introduce such an accumulator in DynamicPPL in the future, but for now you'll need to do it yourself.
+  - `tilde_observe` and `observe` have been removed. `tilde_observe!!` still exists, and any contexts should modify its behaviour. We may further rework the call stack under `tilde_observe!!` in the near future.
+  - `tilde_assume` no longer returns the log density of the current assumption as its second return value. We may further rework the `tilde_assume!!` call stack as well.
+  - For literal observation statements like `0.0 ~ Normal(blahblah)` we used to call `tilde_observe!!` without the `vn` argument. This method no longer exists. Rather we call `tilde_observe!!` with `vn` set to `nothing`.
+  - `set/reset/increment_num_produce!` have become `set/reset/increment_num_produce!!` (note the second exclamation mark). They are no longer guaranteed to modify the `VarInfo` in place, and one should always use the return value.
+  - `@addlogprob!` now _always_ adds to the log likelihood. Previously it added to the log probability that the execution context specified, e.g. the log prior when using `PriorContext`.
+  - `getlogp` now returns a `NamedTuple` with keys `logprior` and `loglikelihood`. If you want the log joint probability, which is what `getlogp` used to return, use `getlogjoint`.
+  - Correspondingly `setlogp!!` and `acclogp!!` should now be called with a `NamedTuple` with keys `logprior` and `loglikelihood`. The `acclogp!!` method with a single scalar value has been deprecated and falls back on `accloglikelihood!!`, and the single scalar version of `setlogp!!` has been removed. Corresponding setter/accumulator functions exist for the log prior as well.
+
 ## 0.36.0
 
 **Breaking changes**

diff --git a/Project.toml b/Project.toml
@@ -21,6 +21,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -68,6 +69,7 @@ MCMCChains = "6"
 MacroTools = "0.5.6"
 Mooncake = "0.4.95"
 OrderedCollections = "1"
+Printf = "1.10"
 Random = "1.6"
 Requires = "1"
 Statistics = "1"

diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl
@@ -100,4 +100,5 @@ PrettyTables.pretty_table(
     header=header,
     tf=PrettyTables.tf_markdown,
     formatters=ft_printf("%.1f", [6, 7]),
+    crop=:none,  # Always print the whole table, even if it doesn't fit in the terminal.
 )
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -160,7 +160,7 @@ returned(::Model)
 
 ## Utilities
 
-It is possible to manually increase (or decrease) the accumulated log density from within a model function.
+It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function.
 
 ```@docs
 @addlogprob!
@@ -328,9 +328,9 @@ The following functions were used for sequential Monte Carlo methods.
 
 ```@docs
 get_num_produce
-set_num_produce!
-increment_num_produce!
-reset_num_produce!
+set_num_produce!!
+increment_num_produce!!
+reset_num_produce!!
 setorder!
 set_retained_vns_del!
 ```
@@ -345,6 +345,22 @@ Base.empty!
 SimpleVarInfo
 ```
 
+### Accumulators
+
+The subtypes of [`AbstractVarInfo`](@ref) store the cumulative log prior and log likelihood, and sometimes other variables that change during executing, in what are called accumulators.
+
+```@docs
+AbstractAccumulator
+```
+
+DynamicPPL provides the following default accumulators.
+
+```@docs
+LogPriorAccumulator
+LogLikelihoodAccumulator
+NumProduceAccumulator
+```
+
 ### Common API
 
 #### Accumulation of log-probabilities
@@ -353,6 +369,13 @@ SimpleVarInfo
 getlogp
 setlogp!!
 acclogp!!
+getlogjoint
+getlogprior
+setlogprior!!
+acclogprior!!
+getloglikelihood
+setloglikelihood!!
+accloglikelihood!!
 resetlogp!!
 ```
 
@@ -427,9 +450,6 @@ Contexts are subtypes of `AbstractPPL.AbstractContext`.
 ```@docs
 SamplingContext
 DefaultContext
-LikelihoodContext
-PriorContext
-MiniBatchContext
 PrefixContext
 ConditionContext
 ```
@@ -476,7 +496,3 @@ DynamicPPL.Experimental.is_suitable_varinfo
 ```@docs
 tilde_assume
 ```
-
-```@docs
-tilde_observe
-```
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
@@ -48,18 +48,18 @@ end
 Sample from the posterior predictive distribution by executing `model` with parameters fixed to each sample
 in `chain`, and return the resulting `Chains`.
 
-The `model` passed to `predict` is often different from the one used to generate `chain`. 
-Typically, the model from which `chain` originated treats certain variables as observed (i.e., 
-data points), while the model you pass to `predict` may mark these same variables as missing 
-or unobserved. Calling `predict` then leverages the previously inferred parameter values to 
+The `model` passed to `predict` is often different from the one used to generate `chain`.
+Typically, the model from which `chain` originated treats certain variables as observed (i.e.,
+data points), while the model you pass to `predict` may mark these same variables as missing
+or unobserved. Calling `predict` then leverages the previously inferred parameter values to
 simulate what new, unobserved data might look like, given your posterior beliefs.
 
 For each parameter configuration in `chain`:
 1. All random variables present in `chain` are fixed to their sampled values.
 2. Any variables not included in `chain` are sampled from their prior distributions.
 
 If `include_all` is `false`, the returned `Chains` will contain only those variables that were not fixed by
-the samples in `chain`. This is useful when you want to sample only new variables from the posterior 
+the samples in `chain`. This is useful when you want to sample only new variables from the posterior
 predictive distribution.
 
 # Examples
@@ -124,7 +124,7 @@ function DynamicPPL.predict(
             map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals)),
         )
 
-        return (varname_and_values=varname_vals, logp=DynamicPPL.getlogp(varinfo))
+        return (varname_and_values=varname_vals, logp=DynamicPPL.getlogjoint(varinfo))
     end
 
     chain_result = reduce(

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
@@ -6,6 +6,7 @@ using Bijectors
 using Compat
 using Distributions
 using OrderedCollections: OrderedCollections, OrderedDict
+using Printf: Printf
 
 using AbstractMCMC: AbstractMCMC
 using ADTypes: ADTypes
@@ -46,17 +47,28 @@ import Base:
 export AbstractVarInfo,
     VarInfo,
     SimpleVarInfo,
+    AbstractAccumulator,
+    LogLikelihoodAccumulator,
+    LogPriorAccumulator,
+    NumProduceAccumulator,
     push!!,
     empty!!,
     subset,
     getlogp,
+    getlogjoint,
+    getlogprior,
+    getloglikelihood,
     setlogp!!,
+    setlogprior!!,
+    setloglikelihood!!,
     acclogp!!,
+    acclogprior!!,
+    accloglikelihood!!,
     resetlogp!!,
     get_num_produce,
-    set_num_produce!,
-    reset_num_produce!,
-    increment_num_produce!,
+    set_num_produce!!,
+    reset_num_produce!!,
+    increment_num_produce!!,
     set_retained_vns_del!,
     is_flagged,
     set_flag!,
@@ -92,15 +104,10 @@ export AbstractVarInfo,
     # Contexts
     SamplingContext,
     DefaultContext,
-    LikelihoodContext,
-    PriorContext,
-    MiniBatchContext,
     PrefixContext,
     ConditionContext,
     assume,
-    observe,
     tilde_assume,
-    tilde_observe,
     # Pseudo distributions
     NamedDist,
     NoDist,
@@ -146,6 +153,9 @@ macro prob_str(str)
     ))
 end
 
+# TODO(mhauru) We should write down the list of methods that any subtype of AbstractVarInfo
+# has to implement. Not sure what the full list is for parameters values, but for
+# accumulators we only need `getaccs` and `setaccs!!`.
 """
     AbstractVarInfo
 
@@ -166,6 +176,8 @@ include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
 include("varnamedvector.jl")
+include("accumulators.jl")
+include("default_accumulators.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")