Accumulators, stage 1

HISTORY.md

@@ -1,5 +1,87 @@
 # DynamicPPL Changelog
 
+## 0.36.0
+
+**Breaking changes**
+
+### VarInfo constructors
+
+`VarInfo(vi::VarInfo, values)` has been removed. You can replace this directly with `unflatten(vi, values)` instead.
+
+The `metadata` argument to `VarInfo([rng, ]model[, sampler, context, metadata])` has been removed.
+If you were not using this argument (most likely), then there is no change needed.
+If you were using the `metadata` argument to specify a blank `VarNamedVector`, then you should replace calls to `VarInfo` with `DynamicPPL.typed_vector_varinfo` instead (see 'Other changes' below).
+
+The `UntypedVarInfo` constructor and type is no longer exported.
+If you needed to construct one, you should now use `DynamicPPL.untyped_varinfo` instead.
+
+The `TypedVarInfo` constructor and type is no longer exported.
+The _type_ has been replaced with `DynamicPPL.NTVarInfo`.
+The _constructor_ has been replaced with `DynamicPPL.typed_varinfo`.
+
+Note that the exact kind of VarInfo returned by `VarInfo(rng, model, ...)` is an implementation detail.
+Previously, it was guaranteed that this would always be a VarInfo whose metadata was a `NamedTuple` containing `Metadata` structs.
+Going forward, this is no longer the case, and you should only assume that the returned object obeys the `AbstractVarInfo` interface.
+
+### VarName prefixing behaviour
+
+The way in which VarNames in submodels are prefixed has been changed.
+This is best explained through an example.
+Consider this model and submodel:
+
+```julia
+using DynamicPPL, Distributions
+@model inner() = x ~ Normal()
+@model outer() = a ~ to_submodel(inner())
+```
+
+In previous versions, the inner variable `x` would be saved as `a.x`.
+However, this was represented as a single symbol `Symbol("a.x")`:
+
+```julia
+julia> dump(keys(VarInfo(outer()))[1])
+VarName{Symbol("a.x"), typeof(identity)}
+  optic: identity (function of type typeof(identity))
+```
+
+Now, the inner variable is stored as a field `x` on the VarName `a`:
+
+```julia
+julia> dump(keys(VarInfo(outer()))[1])
+VarName{:a, Accessors.PropertyLens{:x}}
+  optic: Accessors.PropertyLens{:x} (@o _.x)
+```
+
+In practice, this means that if you are trying to condition a variable in the submodel, you now need to use
+
+```julia
+outer() | (@varname(a.x) => 1.0,)
+```
+
+instead of either of these (which would have worked previously)
+
+```julia
+outer() | (@varname(var"a.x") => 1.0,)
+outer() | (a.x=1.0,)
+```
+
+If you are sampling from a model with submodels, this doesn't affect the way you interact with the `MCMCChains.Chains` object, because VarNames are converted into Symbols when stored in the chain.
+(This behaviour will likely be changed in the future, in that Chains should be indexable by VarNames and not just Symbols, but that has not been implemented yet.)
+
+**Other changes**
+
+While these are technically breaking, they are only internal changes and do not affect the public API.
+The following four functions have been added and/or reworked to make it easier to construct VarInfos with different types of metadata:
+
+ 1. `DynamicPPL.untyped_varinfo([rng, ]model[, sampler, context])`
+ 2. `DynamicPPL.typed_varinfo([rng, ]model[, sampler, context])`
+ 3. `DynamicPPL.untyped_vector_varinfo([rng, ]model[, sampler, context])`
+ 4. `DynamicPPL.typed_vector_varinfo([rng, ]model[, sampler, context])`
+
+The reason for this change is that there were several flavours of VarInfo.
+Some, like `typed_varinfo`, were easy to construct because we had convenience methods for them; however, the others were more difficult.
+This change makes it easier to access different VarInfo types, and also makes it more explicit which one you are constructing.
+
 ## 0.35.6
 
 Fixed the implementation of `.~`, such that running a model with it no longer requires DynamicPPL itself to be loaded.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.35.6"
+version = "0.36.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -44,7 +44,7 @@ DynamicPPLMooncakeExt = ["Mooncake"]
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.10.1"
+AbstractPPL = "0.11"
 Accessors = "0.1"
 BangBang = "0.4.1"
 Bijectors = "0.13.18, 0.14, 0.15"

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -52,8 +52,8 @@ end
 
 Create a benchmark suite for `model` using the selected varinfo type and AD backend.
 Available varinfo choices:
-  • `:untyped`           → uses `VarInfo()`
-  • `:typed`             → uses `VarInfo(model)`
+  • `:untyped`           → uses `DynamicPPL.untyped_varinfo(model)`
+  • `:typed`             → uses `DynamicPPL.typed_varinfo(model)`
   • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())`
   • `:simple_dict`       → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs)
 
@@ -67,11 +67,9 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
     suite = BenchmarkGroup()
 
     vi = if varinfo_choice == :untyped
-        vi = VarInfo()
-        model(rng, vi)
-        vi
+        DynamicPPL.untyped_varinfo(rng, model)
     elseif varinfo_choice == :typed
-        VarInfo(rng, model)
+        DynamicPPL.typed_varinfo(rng, model)
     elseif varinfo_choice == :simple_namedtuple
         SimpleVarInfo{Float64}(model(rng))
     elseif varinfo_choice == :simple_dict

docs/Project.toml

@@ -4,6 +4,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177"
+DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"

docs/src/api.md

@@ -149,7 +149,7 @@ In the past, one would instead embed sub-models using [`@submodel`](@ref), which
 In the context of including models within models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
 
 ```@docs
-prefix
+DynamicPPL.prefix
 ```
 
 Under the hood, [`to_submodel`](@ref) makes use of the following method to indicate that the model it's wrapping is a model over its return-values rather than something else
@@ -291,18 +291,17 @@ AbstractVarInfo
 
 But exactly how a [`AbstractVarInfo`](@ref) stores this information can vary.
 
-For constructing the "default" typed and untyped varinfo types used in DynamicPPL (see [the section on varinfo design](@ref "Design of `VarInfo`") for more on this), we have the following two methods:
+#### `VarInfo`
 
 ```@docs
-DynamicPPL.untyped_varinfo
-DynamicPPL.typed_varinfo
+VarInfo
 ```
 
-#### `VarInfo`
-
 ```@docs
-VarInfo
-TypedVarInfo
+DynamicPPL.untyped_varinfo
+DynamicPPL.typed_varinfo
+DynamicPPL.untyped_vector_varinfo
+DynamicPPL.typed_vector_varinfo
 ```
 
 One main characteristic of [`VarInfo`](@ref) is that samples are transformed to unconstrained Euclidean space and stored in a linearized form, as described in the [main Turing documentation](https://turinglang.org/docs/developers/transforms/dynamicppl/).
@@ -420,7 +419,6 @@ SamplingContext
 DefaultContext
 LikelihoodContext
 PriorContext
-MiniBatchContext
 PrefixContext
 ConditionContext
 ```

docs/src/internals/varinfo.md

@@ -227,13 +227,13 @@ Continuing from the example from the previous section, we can use a `VarInfo` wi
 
 ```@example varinfo-design
 # Type-unstable
-varinfo_untyped_vnv = DynamicPPL.VectorVarInfo(varinfo_untyped)
+varinfo_untyped_vnv = DynamicPPL.untyped_vector_varinfo(varinfo_untyped)
 varinfo_untyped_vnv[@varname(x)], varinfo_untyped_vnv[@varname(y)]
 ```
 
 ```@example varinfo-design
 # Type-stable
-varinfo_typed_vnv = DynamicPPL.VectorVarInfo(varinfo_typed)
+varinfo_typed_vnv = DynamicPPL.typed_vector_varinfo(varinfo_typed)
 varinfo_typed_vnv[@varname(x)], varinfo_typed_vnv[@varname(y)]
 ```
 

src/DynamicPPL.jl

@@ -21,6 +21,9 @@ using DocStringExtensions
 
 using Random: Random
 
+# For extending
+import AbstractPPL: predict
+
 # TODO: Remove these when it's possible.
 import Bijectors: link, invlink
 
@@ -39,13 +42,9 @@ import Base:
     keys,
     haskey
 
-import AbstractPPL: predict
-
 # VarInfo
 export AbstractVarInfo,
     VarInfo,
-    UntypedVarInfo,
-    TypedVarInfo,
     SimpleVarInfo,
     push!!,
     empty!!,
@@ -55,9 +54,9 @@ export AbstractVarInfo,
     acclogp!!,
     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!,
@@ -93,9 +92,6 @@ export AbstractVarInfo,
     # Contexts
     SamplingContext,
     DefaultContext,
-    LikelihoodContext,
-    PriorContext,
-    MiniBatchContext,
     PrefixContext,
     ConditionContext,
     assume,
@@ -167,6 +163,7 @@ include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
 include("varnamedvector.jl")
+include("accumulators.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")

src/abstract_varinfo.jl

@@ -91,36 +91,70 @@ function transformation end
 
 # Accumulation of log-probabilities.
 """
-    getlogp(vi::AbstractVarInfo)
+    getlogjoint(vi::AbstractVarInfo)
 
 Return the log of the joint probability of the observed data and parameters sampled in
 `vi`.
 """
-function getlogp end
+getlogjoint(vi::AbstractVarInfo) = getlogprior(vi) + getloglikelihood(vi)
+getlogp(vi::AbstractVarInfo) = getlogjoint(vi)
+
+function setaccs!! end
+function getaccs end
+
+getlogprior(vi::AbstractVarInfo) = getacc(vi, LogPrior).logp
+getloglikelihood(vi::AbstractVarInfo) = getacc(vi, LogLikelihood).logp
+
+function setacc!!(vi::AbstractVarInfo, acc::AbstractAccumulator)
+    return setaccs!!(vi, setacc!!(getaccs(vi), acc))
+end
+
+setlogprior!!(vi::AbstractVarInfo, logp) = setacc!!(vi, LogPrior(logp))
+setloglikelihood!!(vi::AbstractVarInfo, logp) = setacc!!(vi, LogLikelihood(logp))
 
 """
     setlogp!!(vi::AbstractVarInfo, logp)
 
 Set the log of the joint probability of the observed data and parameters sampled in
 `vi` to `logp`, mutating if it makes sense.
 """
-function setlogp!! end
+function setlogp!!(vi::AbstractVarInfo, logp)
+    vi = setlogprior!!(vi, zero(logp))
+    vi = setloglikelihood!!(vi, logp)
+    return vi
+end
+
+function getacc(vi::AbstractVarInfo, ::Type{AccType}) where {AccType}
+    return getacc(getaccs(vi), AccType)
+end
+
+function accumulate_assume!!(vi::AbstractVarInfo, r, logjac, vn, right)
+    return setaccs!!(vi, accumulate_assume!!(getaccs(vi), r, logjac, vn, right))
+end
+
+function accumulate_observe!!(vi::AbstractVarInfo, left, right)
+    return setaccs!!(vi, accumulate_observe!!(getaccs(vi), left, right))
+end
+
+function acc!!(vi::AbstractVarInfo, ::Type{AccType}, args...) where {AccType}
+    return setaccs!!(vi, acc!!(getaccs(vi), AccType, args...))
+end
+
+function acclogprior!!(vi::AbstractVarInfo, logp)
+    return acc!!(vi, LogPrior, logp)
+end
+
+function accloglikelihood!!(vi::AbstractVarInfo, logp)
+    return acc!!(vi, LogLikelihood, logp)
+end
 
 """
-    acclogp!!([context::AbstractContext, ]vi::AbstractVarInfo, logp)
+    acclogp!!(vi::AbstractVarInfo, logp)
 
 Add `logp` to the value of the log of the joint probability of the observed data and
 parameters sampled in `vi`, mutating if it makes sense.
 """
-function acclogp!!(context::AbstractContext, vi::AbstractVarInfo, logp)
-    return acclogp!!(NodeTrait(context), context, vi, logp)
-end
-function acclogp!!(::IsLeaf, context::AbstractContext, vi::AbstractVarInfo, logp)
-    return acclogp!!(vi, logp)
-end
-function acclogp!!(::IsParent, context::AbstractContext, vi::AbstractVarInfo, logp)
-    return acclogp!!(childcontext(context), vi, logp)
-end
+acclogp!!(vi::AbstractVarInfo, logp) = accloglikelihood!!(vi, logp)
 
 """
     resetlogp!!(vi::AbstractVarInfo)
@@ -247,11 +281,11 @@ julia> values_as(SimpleVarInfo(data), Vector)
  2.0
 ```
 
-`TypedVarInfo`:
+`VarInfo` with `NamedTuple` of `Metadata`:
 
 ```jldoctest
 julia> # Just use an example model to construct the `VarInfo` because we're lazy.
-       vi = VarInfo(DynamicPPL.TestUtils.demo_assume_dot_observe());
+       vi = DynamicPPL.typed_varinfo(DynamicPPL.TestUtils.demo_assume_dot_observe());
 
 julia> vi[@varname(s)] = 1.0; vi[@varname(m)] = 2.0;
 
@@ -273,11 +307,11 @@ julia> values_as(vi, Vector)
  2.0
 ```
 
-`UntypedVarInfo`:
+`VarInfo` with `Metadata`:
 
 ```jldoctest
 julia> # Just use an example model to construct the `VarInfo` because we're lazy.
-       vi = VarInfo(); DynamicPPL.TestUtils.demo_assume_dot_observe()(vi);
+       vi = DynamicPPL.untyped_varinfo(DynamicPPL.TestUtils.demo_assume_dot_observe());
 
 julia> vi[@varname(s)] = 1.0; vi[@varname(m)] = 2.0;
 
@@ -725,7 +759,7 @@ end
 
 # Legacy code that is currently overloaded for the sake of simplicity.
 # TODO: Remove when possible.
-increment_num_produce!(::AbstractVarInfo) = nothing
+increment_num_produce!!(::AbstractVarInfo) = nothing
 
 """
     from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])