Remove SimpleVarInfo

Author: penelopeysm

HISTORY.md

@@ -15,6 +15,9 @@ As a result of this change, `LogDensityFunction` no longer stores a VarInfo insi
 In general, if `ldf` is a `LogDensityFunction`, it is now only valid to access `ldf.model` and `ldf.adtype`.
 If you were previously relying on this behaviour, you will need to store a VarInfo separately.
 
+Along with this change, DynamicPPL now exposes the `fast_evaluate!!` method which allows you to hook into this 'fast evaluation' pipeline directly.
+Please see the documentation for details.
+
 #### Parent and leaf contexts
 
 The `DynamicPPL.NodeTrait` function has been removed.
@@ -28,6 +31,17 @@ Leaf contexts require no changes, apart from a removal of the `NodeTrait` functi
 `ConditionContext` and `PrefixContext` are no longer exported.
 You should not need to use these directly, please use `AbstractPPL.condition` and `DynamicPPL.prefix` instead.
 
+#### SimpleVarInfo
+
+`SimpleVarInfo` has been removed.
+Its main purpose was for evaluating models rapidly.
+However, `fast_evaluate!!` provides a cleaner way of doing this.
+In particular, if you want to evaluate a model at a given set of parameters, you can do:
+
+```julia
+retval, vi = DynamicPPL.fast_evaluate!!(rng, model, InitFromParams(params), accs)
+```
+
 #### Miscellaneous
 
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.

benchmarks/benchmarks.jl

@@ -59,9 +59,7 @@ chosen_combinations = [
         false,
     ),
     ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff, false),
-    ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true),
-    ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed_vector, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :untyped_vector, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -1,6 +1,6 @@
 module DynamicPPLBenchmarks
 
-using DynamicPPL: VarInfo, SimpleVarInfo, VarName
+using DynamicPPL: VarInfo, VarName
 using BenchmarkTools: BenchmarkGroup, @benchmarkable
 using DynamicPPL: DynamicPPL
 using ADTypes: ADTypes
@@ -58,8 +58,6 @@ Create a benchmark suite for `model` using the selected varinfo type and AD back
 Available varinfo choices:
   • `: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)
 
 The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversediff`, `:zygote`).
 
@@ -74,12 +72,6 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
         DynamicPPL.untyped_varinfo(rng, model)
     elseif varinfo_choice == :typed
         DynamicPPL.typed_varinfo(rng, model)
-    elseif varinfo_choice == :simple_namedtuple
-        SimpleVarInfo{Float64}(model(rng))
-    elseif varinfo_choice == :simple_dict
-        retvals = model(rng)
-        vns = [VarName{k}() for k in keys(retvals)]
-        SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals))))
     elseif varinfo_choice == :typed_vector
         DynamicPPL.typed_vector_varinfo(rng, model)
     elseif varinfo_choice == :untyped_vector

benchmarks/src/Models.jl

@@ -2,7 +2,7 @@
 Models for benchmarking Turing.jl.
 
 Each model returns a NamedTuple of all the random variables in the model that are not
-observed (this is used for constructing SimpleVarInfos).
+observed.
 """
 module Models
 

docs/src/api.md

@@ -359,12 +359,6 @@ set_transformed!!
 Base.empty!
 ```
 
-#### `SimpleVarInfo`
-
-```@docs
-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.

src/DynamicPPL.jl

@@ -46,7 +46,6 @@ import Base:
 # VarInfo
 export AbstractVarInfo,
     VarInfo,
-    SimpleVarInfo,
     AbstractAccumulator,
     LogLikelihoodAccumulator,
     LogPriorAccumulator,
@@ -174,7 +173,7 @@ Abstract supertype for data structures that capture random variables when execut
 probabilistic model and accumulate log densities such as the log likelihood or the
 log joint probability of the model.
 
-See also: [`VarInfo`](@ref), [`SimpleVarInfo`](@ref).
+See also: [`VarInfo`](@ref).
 """
 abstract type AbstractVarInfo <: AbstractModelTrace end
 
@@ -196,7 +195,6 @@ include("default_accumulators.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")
-include("simple_varinfo.jl")
 include("onlyaccs.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")

src/abstract_varinfo.jl

@@ -502,52 +502,6 @@ If no `Type` is provided, return values as stored in `varinfo`.
 
 # Examples
 
-`SimpleVarInfo` with `NamedTuple`:
-
-```jldoctest
-julia> data = (x = 1.0, m = [2.0]);
-
-julia> values_as(SimpleVarInfo(data))
-(x = 1.0, m = [2.0])
-
-julia> values_as(SimpleVarInfo(data), NamedTuple)
-(x = 1.0, m = [2.0])
-
-julia> values_as(SimpleVarInfo(data), OrderedDict)
-OrderedDict{VarName{sym, typeof(identity)} where sym, Any} with 2 entries:
-  x => 1.0
-  m => [2.0]
-
-julia> values_as(SimpleVarInfo(data), Vector)
-2-element Vector{Float64}:
- 1.0
- 2.0
-```
-
-`SimpleVarInfo` with `OrderedDict`:
-
-```jldoctest
-julia> data = OrderedDict{Any,Any}(@varname(x) => 1.0, @varname(m) => [2.0]);
-
-julia> values_as(SimpleVarInfo(data))
-OrderedDict{Any, Any} with 2 entries:
-  x => 1.0
-  m => [2.0]
-
-julia> values_as(SimpleVarInfo(data), NamedTuple)
-(x = 1.0, m = [2.0])
-
-julia> values_as(SimpleVarInfo(data), OrderedDict)
-OrderedDict{Any, Any} with 2 entries:
-  x => 1.0
-  m => [2.0]
-
-julia> values_as(SimpleVarInfo(data), Vector)
-2-element Vector{Float64}:
- 1.0
- 2.0
-```
-
 `VarInfo` with `NamedTuple` of `Metadata`:
 
 ```jldoctest
@@ -828,8 +782,8 @@ function link!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return link!!(default_transformation(model, vi), vi, vns, model)
 end
 function link!!(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
-    # Note that in practice this method is only called for SimpleVarInfo, because VarInfo
-    # has a dedicated implementation
+    # Note that VarInfo has a dedicated implementation so this is only a generic
+    # fallback (previously used for SimpleVarInfo)
     model = setleafcontext(model, DynamicTransformationContext{false}())
     vi = last(evaluate!!(model, vi))
     return set_transformed!!(vi, t)
@@ -890,8 +844,8 @@ function invlink!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return invlink!!(default_transformation(model, vi), vi, vns, model)
 end
 function invlink!!(::DynamicTransformation, vi::AbstractVarInfo, model::Model)
-    # Note that in practice this method is only called for SimpleVarInfo, because VarInfo
-    # has a dedicated implementation
+    # Note that VarInfo has a dedicated implementation so this is only a generic
+    # fallback (previously used for SimpleVarInfo)
     model = setleafcontext(model, DynamicTransformationContext{true}())
     vi = last(evaluate!!(model, vi))
     return set_transformed!!(vi, NoTransformation())
@@ -946,47 +900,6 @@ This will be called prior to `model` evaluation, allowing one to perform a singl
 basis as is done with [`DynamicTransformation`](@ref).
 
 See also: [`StaticTransformation`](@ref), [`DynamicTransformation`](@ref).
-
-# Examples
-```julia-repl
-julia> using DynamicPPL, Distributions, Bijectors
-
-julia> @model demo() = x ~ Normal()
-demo (generic function with 2 methods)
-
-julia> # By subtyping `Transform`, we inherit the `(inv)link!!`.
-       struct MyBijector <: Bijectors.Transform end
-
-julia> # Define some dummy `inverse` which will be used in the `link!!` call.
-       Bijectors.inverse(f::MyBijector) = identity
-
-julia> # We need to define `with_logabsdet_jacobian` for `MyBijector`
-       # (`identity` already has `with_logabsdet_jacobian` defined)
-       function Bijectors.with_logabsdet_jacobian(::MyBijector, x)
-           # Just using a large number of the logabsdet-jacobian term
-           # for demonstration purposes.
-           return (x, 1000)
-       end
-
-julia> # Change the `default_transformation` for our model to be a
-       # `StaticTransformation` using `MyBijector`.
-       function DynamicPPL.default_transformation(::Model{typeof(demo)})
-           return DynamicPPL.StaticTransformation(MyBijector())
-       end
-
-julia> model = demo();
-
-julia> vi = SimpleVarInfo(x=1.0)
-SimpleVarInfo((x = 1.0,), 0.0)
-
-julia> # Uses the `inverse` of `MyBijector`, which we have defined as `identity`
-       vi_linked = link!!(vi, model)
-Transformed SimpleVarInfo((x = 1.0,), 0.0)
-
-julia> # Now performs a single `invlink!!` before model evaluation.
-       logjoint(model, vi_linked)
--1001.4189385332047
-```
 """
 function maybe_invlink_before_eval!!(vi::AbstractVarInfo, model::Model)
     return maybe_invlink_before_eval!!(transformation(vi), vi, model)

src/contexts/transformation.jl

@@ -7,7 +7,7 @@ constrained space if `isinverse` or unconstrained if `!isinverse`.
 Note that some `AbstractVarInfo` types, must notably `VarInfo`, override the
 `DynamicTransformationContext` methods with more efficient implementations.
 `DynamicTransformationContext` is a fallback for when we need to evaluate the model to know
-how to do the transformation, used by e.g. `SimpleVarInfo`.
+how to do the transformation.
 """
 struct DynamicTransformationContext{isinverse} <: AbstractContext end
 

src/fasteval.jl

@@ -334,11 +334,6 @@ tweak_adtype(adtype::ADTypes.AbstractADType, ::Model, ::AbstractVarInfo) = adtyp
 # Helper functions to extract ranges and link status #
 ######################################################
 
-# This fails for SimpleVarInfo, but honestly there is no reason to support that here. The
-# fact is that evaluation doesn't use a VarInfo, it only uses it once to generate the ranges
-# and link status. So there is no motivation to use SimpleVarInfo inside a
-# LogDensityFunction any more, we can just always use typed VarInfo. In fact one could argue
-# that there is no purpose in supporting untyped VarInfo either.
 """
     get_ranges_and_linked(varinfo::VarInfo)
 

src/model.jl

@@ -1062,8 +1062,11 @@ Base.nameof(model::Model{<:Function}) = nameof(model.f)
 Generate a sample of type `T` from the prior distribution of the `model`.
 """
 function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
-    x = last(init!!(rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())))
-    return values_as(x, T)
+    # TODO(penelopeysm): This can be done with an accumulator instead. For
+    # T = Dict, ValuesAsInModelAcc can already do it. For T = NamedTuple we
+    # would just need a similar accumulator that collects into a NamedTuple
+    # rather than a Dict.
+    return values_as(VarInfo(rng, model), T)
 end
 
 # Default RNG and type
@@ -1155,12 +1158,115 @@ julia> returned(model, Dict{VarName,Float64}(@varname(m) => 2.0))
 ```
 """
 function returned(model::Model, parameters::Union{NamedTuple,AbstractDict{<:VarName}})
-    vi = DynamicPPL.setaccs!!(VarInfo(), ())
-    # Note: we can't use `fix(model, parameters)` because
-    # https://github.com/TuringLang/DynamicPPL.jl/issues/1097
-    # Use `nothing` as the fallback to ensure that any missing parameters cause an error
-    ctx = InitContext(Random.default_rng(), InitFromParams(parameters, nothing))
-    new_model = setleafcontext(model, ctx)
-    # We can't use new_model() because that overwrites it with an InitContext of its own.
-    return first(evaluate!!(new_model, vi))
+    accs = AccumulatorTuple()
+    retval, _ = DynamicPPL.fast_evaluate!!(model, InitFromParams(parameters, nothing), accs)
+    return retval
+end
+
+"""
+    logjoint(model::Model, θ::Union{NamedTuple,AbstractDict})
+
+Return the log joint probability of variables `θ` for the probabilistic `model`.
+
+See [`logprior`](@ref) and [`loglikelihood`](@ref).
+
+# Examples
+```jldoctest; setup=:(using Distributions)
+julia> @model function demo(x)
+           m ~ Normal()
+           for i in eachindex(x)
+               x[i] ~ Normal(m, 1.0)
+           end
+       end
+demo (generic function with 2 methods)
+
+julia> # Using a `NamedTuple`.
+       logjoint(demo([1.0]), (m = 100.0, ))
+-9902.33787706641
+
+julia> # Using a `OrderedDict`.
+       logjoint(demo([1.0]), OrderedDict(@varname(m) => 100.0))
+-9902.33787706641
+
+julia> # Truth.
+       logpdf(Normal(100.0, 1.0), 1.0) + logpdf(Normal(), 100.0)
+-9902.33787706641
+```
+"""
+function logjoint(model::Model, θ::Union{NamedTuple,AbstractDict})
+    accs = AccumulatorTuple((LogPriorAccumulator(), LogLikelihoodAccumulator()))
+    _, vi = DynamicPPL.fast_evaluate!!(model, InitFromParams(θ, nothing), accs)
+    return getlogjoint(vi)
+end
+
+"""
+    logprior(model::Model, θ::Union{NamedTuple,AbstractDict})
+
+Return the log prior probability of variables `θ` for the probabilistic `model`.
+
+See also [`logjoint`](@ref) and [`loglikelihood`](@ref).
+
+# Examples
+```jldoctest; setup=:(using Distributions)
+julia> @model function demo(x)
+           m ~ Normal()
+           for i in eachindex(x)
+               x[i] ~ Normal(m, 1.0)
+           end
+       end
+demo (generic function with 2 methods)
+
+julia> # Using a `NamedTuple`.
+       logprior(demo([1.0]), (m = 100.0, ))
+-5000.918938533205
+
+julia> # Using a `OrderedDict`.
+       logprior(demo([1.0]), OrderedDict(@varname(m) => 100.0))
+-5000.918938533205
+
+julia> # Truth.
+       logpdf(Normal(), 100.0)
+-5000.918938533205
+```
+"""
+function logprior(model::Model, θ::Union{NamedTuple,AbstractDict})
+    accs = AccumulatorTuple((LogPriorAccumulator(),))
+    _, vi = DynamicPPL.fast_evaluate!!(model, InitFromParams(θ, nothing), accs)
+    return getlogprior(vi)
+end
+
+"""
+    loglikelihood(model::Model, θ::Union{NamedTuple,AbstractDict})
+
+Return the log likelihood of variables `θ` for the probabilistic `model`.
+
+See also [`logjoint`](@ref) and [`logprior`](@ref).
+
+# Examples
+```jldoctest; setup=:(using Distributions)
+julia> @model function demo(x)
+           m ~ Normal()
+           for i in eachindex(x)
+               x[i] ~ Normal(m, 1.0)
+           end
+       end
+demo (generic function with 2 methods)
+
+julia> # Using a `NamedTuple`.
+       loglikelihood(demo([1.0]), (m = 100.0, ))
+-4901.418938533205
+
+julia> # Using a `OrderedDict`.
+       loglikelihood(demo([1.0]), OrderedDict(@varname(m) => 100.0))
+-4901.418938533205
+
+julia> # Truth.
+       logpdf(Normal(100.0, 1.0), 1.0)
+-4901.418938533205
+```
+"""
+function Distributions.loglikelihood(model::Model, θ::Union{NamedTuple,AbstractDict})
+    accs = AccumulatorTuple((LogLikelihoodAccumulator(),))
+    _, vi = DynamicPPL.fast_evaluate!!(model, InitFromParams(θ, nothing), accs)
+    return getloglikelihood(vi)
 end