Merge branch 'breaking' into mhauru/accumulators-stage2

Author: penelopeysm

HISTORY.md

@@ -2,12 +2,57 @@
 
 ## 0.37.0
 
-**Breaking changes**
+DynamicPPL 0.37 comes with a substantial reworking of its internals.
+Fundamentally, there is no change to the actual modelling syntax: if you are a Turing.jl user, for example, this release is unlikely to affect you much.
+However, if you are a package developer or someone who uses DynamicPPL's functionality directly, you will notice a number of changes.
+
+To avoid overwhelming the reader, we begin by listing the most important, user-facing changes, before explaining the changes to the internals in more detail.
+
+Note that virtually all changes listed here are breaking.
+
+**Public-facing changes**
 
 ### Submodel macro
 
 The `@submodel` macro is fully removed; please use `to_submodel` instead.
 
+### `DynamicPPL.TestUtils.AD.run_ad`
+
+The three keyword arguments, `test`, `reference_backend`, and `expected_value_and_grad` have been merged into a single `test` keyword argument.
+Please see the API documentation for more details.
+(The old `test=true` and `test=false` values are still valid, and you only need to adjust the invocation if you were explicitly passing the `reference_backend` or `expected_value_and_grad` arguments.)
+
+There is now also an `rng` keyword argument to help seed parameter generation.
+
+Finally, instead of specifying `value_atol` and `grad_atol`, you can now specify `atol` and `rtol` which are used for both value and gradient.
+Their semantics are the same as in Julia's `isapprox`; two values are equal if they satisfy either `atol` or `rtol`.
+
+### `DynamicPPL.TestUtils.check_model`
+
+You now need to explicitly pass a `VarInfo` argument to `check_model` and `check_model_and_trace`.
+Previously, these functions would generate a new VarInfo for you (using an optionally provided `rng`).
+
+### Removal of `PriorContext` and `LikelihoodContext`
+
+A number of DynamicPPL's contexts have been removed, most notably `PriorContext` and `LikelihoodContext`.
+Although these are not the only _exported_ contexts, we consider unlikely that anyone was using _other_ contexts manually: if you have a question about contexts _other_ than these, please continue reading the 'Internals' section below.
+
+Previously, during evaluation of a model, DynamicPPL only had the capability to store a _single_ log probability (`logp`) field.
+`DefaultContext`, `PriorContext`, and `LikelihoodContext` were used to control what this field represented: they would accumulate the log joint, log prior, or log likelihood, respectively.
+
+Now, we have reworked DynamicPPL's `VarInfo` object such that it can track multiple log probabilities at once (see the 'Accumulators' section below).
+If you were evaluating a model with `PriorContext`, you can now just evaluate it with `DefaultContext`, and instead of calling `getlogp(varinfo)`, you can call `getlogprior(varinfo)` (and similarly for the likelihood).
+
+If you were constructing a `LogDensityFunction` with `PriorContext`, you can now stick to `DefaultContext`.
+`LogDensityFunction` now has an extra field, called `getlogdensity`, which represents a function that takes a `VarInfo` and returns the log density you want.
+Thus, if you pass `getlogprior` as the value of this parameter, you will get the same behaviour as with `PriorContext`.
+
+The other case where one might use `PriorContext` was to use `@addlogprob!` to add to the log prior.
+Previously, this was accomplished by manually checking `__context__ isa DynamicPPL.PriorContext`.
+Now, you can write `@addlogprob (; logprior=x, loglikelihood=y)` to add `x` to the log-prior and `y` to the log-likelihood.
+
+**Internals**
+
 ### 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:
@@ -59,6 +104,18 @@ And a couple of more internal changes:
   - The model evaluation function, `model.f` for some `model::Model`, no longer takes a context as an argument
   - The internal representation and API dealing with submodels (i.e., `ReturnedModelWrapper`, `Sampleable`, `should_auto_prefix`, `is_rhs_model`) has been simplified. If you need to check whether something is a submodel, just use `x isa DynamicPPL.Submodel`. Note that the public API i.e. `to_submodel` remains completely untouched.
 
+## 0.36.15
+
+Bumped minimum Julia version to 1.10.8 to avoid potential crashes with `Core.Compiler.widenconst` (which Mooncake uses).
+
+## 0.36.14
+
+Added compatibility with [email protected].
+
+## 0.36.13
+
+Added documentation for the `returned(::Model, ::MCMCChains.Chains)` method.
+
 ## 0.36.12
 
 Removed several unexported functions.

Project.toml

@@ -47,7 +47,7 @@ DynamicPPLMooncakeExt = ["Mooncake"]
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.11"
+AbstractPPL = "0.11, 0.12"
 Accessors = "0.1"
 BangBang = "0.4.1"
 Bijectors = "0.13.18, 0.14, 0.15"
@@ -74,4 +74,4 @@ Random = "1.6"
 Requires = "1"
 Statistics = "1"
 Test = "1.6"
-julia = "1.10"
+julia = "1.10.8"

docs/Project.toml

@@ -14,7 +14,7 @@ MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
 [compat]
-AbstractPPL = "0.11"
+AbstractPPL = "0.11, 0.12"
 Accessors = "0.1"
 DataStructures = "0.18"
 Distributions = "0.25"

docs/src/api.md

@@ -160,9 +160,10 @@ It is possible to manually increase (or decrease) the accumulated log likelihood
 @addlogprob!
 ```
 
-Return values of the model function for a collection of samples can be obtained with [`returned(model, chain)`](@ref).
+Return values of the model function can be obtained with [`returned(model, sample)`](@ref), where `sample` is either a `MCMCChains.Chains` object (which represents a collection of samples) or a single sample represented as a `NamedTuple`.
 
 ```@docs
+returned(::DynamicPPL.Model, ::MCMCChains.Chains)
 returned(::DynamicPPL.Model, ::NamedTuple)
 ```
 
@@ -205,6 +206,21 @@ To test and/or benchmark the performance of an AD backend on a model, DynamicPPL
 
 ```@docs
 DynamicPPL.TestUtils.AD.run_ad
+```
+
+The default test setting is to compare against ForwardDiff.
+You can have more fine-grained control over how to test the AD backend using the following types:
+
+```@docs
+DynamicPPL.TestUtils.AD.AbstractADCorrectnessTestSetting
+DynamicPPL.TestUtils.AD.WithBackend
+DynamicPPL.TestUtils.AD.WithExpectedResult
+DynamicPPL.TestUtils.AD.NoTest
+```
+
+These are returned / thrown by the `run_ad` function:
+
+```@docs
 DynamicPPL.TestUtils.AD.ADResult
 DynamicPPL.TestUtils.AD.ADIncorrectException
 ```
@@ -325,7 +341,7 @@ get_num_produce
 set_num_produce!!
 increment_num_produce!!
 reset_num_produce!!
-setorder!
+setorder!!
 set_retained_vns_del!
 ```
 
@@ -352,7 +368,7 @@ DynamicPPL provides the following default accumulators.
 ```@docs
 LogPriorAccumulator
 LogLikelihoodAccumulator
-NumProduceAccumulator
+VariableOrderAccumulator
 ```
 
 ### Common API

src/DynamicPPL.jl

@@ -50,7 +50,7 @@ export AbstractVarInfo,
     AbstractAccumulator,
     LogLikelihoodAccumulator,
     LogPriorAccumulator,
-    NumProduceAccumulator,
+    VariableOrderAccumulator,
     push!!,
     empty!!,
     subset,
@@ -73,7 +73,7 @@ export AbstractVarInfo,
     is_flagged,
     set_flag!,
     unset_flag!,
-    setorder!,
+    setorder!!,
     istrans,
     link,
     link!!,

src/abstract_varinfo.jl

@@ -374,6 +374,24 @@ function resetlogp!!(vi::AbstractVarInfo)
     return vi
 end
 
+"""
+    setorder!!(vi::AbstractVarInfo, vn::VarName, index::Integer)
+
+Set the `order` of `vn` in `vi` to `index`, where `order` is the number of `observe
+statements run before sampling `vn`.
+"""
+function setorder!!(vi::AbstractVarInfo, vn::VarName, index::Integer)
+    return map_accumulator!!(acc -> (acc.order[vn] = index; acc), vi, Val(:VariableOrder))
+end
+
+"""
+    getorder(vi::VarInfo, vn::VarName)
+
+Get the `order` of `vn` in `vi`, where `order` is the number of `observe` statements
+run before sampling `vn`.
+"""
+getorder(vi::AbstractVarInfo, vn::VarName) = getacc(vi, Val(:VariableOrder)).order[vn]
+
 # Variables and their realizations.
 @doc """
     keys(vi::AbstractVarInfo)
@@ -725,7 +743,15 @@ If `vns` is provided, then only check if this/these varname(s) are transformed.
 """
 istrans(vi::AbstractVarInfo) = istrans(vi, collect(keys(vi)))
 function istrans(vi::AbstractVarInfo, vns::AbstractVector)
-    return !isempty(vns) && all(Base.Fix1(istrans, vi), vns)
+    # This used to be: `!isempty(vns) && all(Base.Fix1(istrans, vi), vns)`.
+    # In theory that should work perfectly fine. For unbeknownst reasons,
+    # Julia 1.10 fails to infer its return type correctly. Thus we use this
+    # slightly longer definition.
+    isempty(vns) && return false
+    for vn in vns
+        istrans(vi, vn) || return false
+    end
+    return true
 end
 
 """
@@ -972,29 +998,37 @@ end
 
 Return the `num_produce` of `vi`.
 """
-get_num_produce(vi::AbstractVarInfo) = getacc(vi, Val(:NumProduce)).num
+get_num_produce(vi::AbstractVarInfo) = getacc(vi, Val(:VariableOrder)).num_produce
 
 """
     set_num_produce!!(vi::AbstractVarInfo, n::Int)
 
 Set the `num_produce` field of `vi` to `n`.
 """
-set_num_produce!!(vi::AbstractVarInfo, n::Int) = setacc!!(vi, NumProduceAccumulator(n))
+function set_num_produce!!(vi::AbstractVarInfo, n::Integer)
+    if hasacc(vi, Val(:VariableOrder))
+        acc = getacc(vi, Val(:VariableOrder))
+        acc = VariableOrderAccumulator(n, acc.order)
+    else
+        acc = VariableOrderAccumulator(n)
+    end
+    return setacc!!(vi, acc)
+end
 
 """
     increment_num_produce!!(vi::AbstractVarInfo)
 
 Add 1 to `num_produce` in `vi`.
 """
 increment_num_produce!!(vi::AbstractVarInfo) =
-    map_accumulator!!(increment, vi, Val(:NumProduce))
+    map_accumulator!!(increment, vi, Val(:VariableOrder))
 
 """
     reset_num_produce!!(vi::AbstractVarInfo)
 
 Reset the value of `num_produce` in `vi` to 0.
 """
-reset_num_produce!!(vi::AbstractVarInfo) = map_accumulator!!(zero, vi, Val(:NumProduce))
+reset_num_produce!!(vi::AbstractVarInfo) = set_num_produce!!(vi, zero(get_num_produce(vi)))
 
 """
     from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

src/accumulators.jl

@@ -13,6 +13,7 @@ An accumulator type `T <: AbstractAccumulator` must implement the following meth
 - `accumulator_name(acc::T)` or `accumulator_name(::Type{T})`
 - `accumulate_observe!!(acc::T, right, left, vn)`
 - `accumulate_assume!!(acc::T, val, logjac, vn, right)`
+- `Base.copy(acc::T)`
 
 To be able to work with multi-threading, it should also implement:
 - `split(acc::T)`
@@ -53,10 +54,11 @@ function accumulate_observe!! end
 
 Update `acc` in a `tilde_assume!!` call. Returns the updated `acc`.
 
-`vn` is the name of the variable being assumed, `val` is the value of the variable, and
-`right` is the distribution on the RHS of the tilde statement. `logjac` is the log
-determinant of the Jacobian of the transformation that was done to convert the value of `vn`
-as it was given (e.g. by sampler operating in linked space) to `val`.
+`vn` is the name of the variable being assumed, `val` is the value of the variable (in the
+original, unlinked space), and `right` is the distribution on the RHS of the tilde
+statement. `logjac` is the log determinant of the Jacobian of the transformation that was
+done to convert the value of `vn` as it was given to `val`: for example, if the sampler is
+operating in linked (Euclidean) space, then logjac will be nonzero.
 
 `accumulate_assume!!` may mutate `acc`, but not any of the other arguments.
 
@@ -71,7 +73,7 @@ Return a new accumulator like `acc` but empty.
 
 The precise meaning of "empty" is that that the returned value should be such that
 `combine(acc, split(acc))` is equal to `acc`. This is used in the context of multi-threading
-where different threads may accumulate independently and the results are the combined.
+where different threads may accumulate independently and the results are then combined.
 
 See also: [`combine`](@ref)
 """
@@ -80,7 +82,8 @@ function split end
 """
     combine(acc::AbstractAccumulator, acc2::AbstractAccumulator)
 
-Combine two accumulators of the same type. Returns a new accumulator.
+Combine two accumulators which have the same type (but may, in general, have different type
+parameters). Returns a new accumulator of the same type.
 
 See also: [`split`](@ref)
 """
@@ -136,6 +139,9 @@ function Base.haskey(at::AccumulatorTuple, ::Val{accname}) where {accname}
     @inline return haskey(at.nt, accname)
 end
 Base.keys(at::AccumulatorTuple) = keys(at.nt)
+Base.:(==)(at1::AccumulatorTuple, at2::AccumulatorTuple) = at1.nt == at2.nt
+Base.hash(at::AccumulatorTuple, h::UInt) = Base.hash((AccumulatorTuple, at.nt), h)
+Base.copy(at::AccumulatorTuple) = AccumulatorTuple(map(copy, at.nt))
 
 function Base.convert(::Type{AccumulatorTuple{N,T}}, accs::AccumulatorTuple{N}) where {N,T}
     return AccumulatorTuple(convert(T, accs.nt))

src/context_implementations.jl

@@ -148,7 +148,6 @@ function assume(
             f = to_maybe_linked_internal_transform(vi, vn, dist)
             # TODO(mhauru) This should probably be call a function called setindex_internal!
             vi = BangBang.setindex!!(vi, f(r), vn)
-            setorder!(vi, vn, get_num_produce(vi))
         else
             # Otherwise we just extract it.
             r = vi[vn, dist]