Remove samplers from VarInfo - indexing (#793)

* Remove selector stuff from varinfo tests

* Implement link and invlink for varnames rather than samplers

* Replace set_retained_vns_del_by_spl! with set_retained_vns_del!

* Make linking tests more extensive

* Remove sampler indexing from link methods (but not invlink)

* Remove indexing by samplers from invlink

* Work towards removing sampler indexing with StaticTransformation

* Fix invlink/link for TypedVarInfo and StaticTransformation

* Fix a test in models.jl

* Move some functions to utils.jl, add tests and docstrings

* Fix a docstring typo

* Various simplification to link/invlink

* Improve a docstring

* Style improvements

* Fix broken link/invlink dispatch cascade for VectorVarInfo

* Fix some more broken dispatch cascades

* Apply suggestions from code review

Co-authored-by: Xianda Sun <[email protected]>

* Remove comments that messed with docstrings

* Apply suggestions from code review

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

* Fix issues surfaced in code review

* Simplify link/invlink arguments

* Fix a bug in unflatten VarNamedVector

* Rename VarNameCollection -> VarNameTuple

* Remove test of a removed varname_namedtuple method

* Apply suggestions from code review

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

* Respond to review feedback

* Remove _default_sampler and a dead argument of maybe_invlink_before_eval

* Fix a typo in a comment

* Add HISTORY entry, fix one set_retained_vns_del! method

* Remove some VarInfo getindex with samplers stuff

* Remove some index setting with samplers

* Remove more sampler indexing

* Remove unflatten with samplers

* Clean up some setindex stuff

* Remove a bunch of varinfo.jl internal functions that used samplers/space, update HISTORY.md

* Fix HISTORY.md

* Miscalleanous small fixes

* Fix a bug in VarInfo constructor

* Fix getparams(::LogDensityFunction)

* Apply suggestions from code review

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

---------

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

Author: 3 people

HISTORY.md

@@ -10,6 +10,11 @@ This release removes the feature of `VarInfo` where it kept track of which varia
 
   - `link` and `invlink`, and their `!!` versions, no longer accept a sampler as an argument to specify which variables to (inv)link. The `link(varinfo, model)` methods remain in place, and as a new addition one can give a `Tuple` of `VarName`s to (inv)link only select variables, as in `link(varinfo, varname_tuple, model)`.
   - `set_retained_vns_del_by_spl!` has been replaced by `set_retained_vns_del!` which applies to all variables.
+  - `getindex`, `setindex!`, and `setindex!!` no longer accept samplers as arguments
+  - `unflatten` no longer accepts a sampler as an argument
+  - `eltype(::VarInfo)` no longer accepts a sampler as an argument
+  - `keys(::VarInfo)` no longer accepts a sampler as an argument
+  - `VarInfo(::VarInfo, ::Sampler, ::AbstactVector)` no longer accepts the sampler argument.
 
 ### Reverse prefixing order
 

src/abstract_varinfo.jl

@@ -149,7 +149,6 @@ If `dist` is specified, the value(s) will be massaged into the representation ex
 
 """
     getindex(vi::AbstractVarInfo, ::Colon)
-    getindex(vi::AbstractVarInfo, ::AbstractSampler)
 
 Return the current value(s) of `vn` (`vns`) in `vi` in the support of its (their)
 distribution(s) as a flattened `Vector`.
@@ -159,7 +158,6 @@ The default implementation is to call [`values_as`](@ref) with `Vector` as the t
 See also: [`getindex(vi::AbstractVarInfo, vn::VarName, dist::Distribution)`](@ref)
 """
 Base.getindex(vi::AbstractVarInfo, ::Colon) = values_as(vi, Vector)
-Base.getindex(vi::AbstractVarInfo, ::AbstractSampler) = vi[:]
 
 """
     getindex_internal(vi::AbstractVarInfo, vn::VarName)
@@ -341,9 +339,9 @@ julia> values_as(vi, Vector)
 function values_as end
 
 """
-    eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromPrior}
+    eltype(vi::AbstractVarInfo)
 
-Determine the default `eltype` of the values returned by `vi[spl]`.
+Return the `eltype` of the values returned by `vi[:]`.
 
 !!! warning
     This should generally not be called explicitly, as it's only used in
@@ -352,13 +350,13 @@ Determine the default `eltype` of the values returned by `vi[spl]`.
 
     This method is considered legacy, and is likely to be deprecated in the future.
 """
-function Base.eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromPrior})
-    T = Base.promote_op(getindex, typeof(vi), typeof(spl))
+function Base.eltype(vi::AbstractVarInfo)
+    T = Base.promote_op(getindex, typeof(vi), Colon)
     if T === Union{}
-        # In this case `getindex(vi, spl)` errors
+        # In this case `getindex(vi, :)` errors
         # Let us throw a more descriptive error message
         # Ref https://github.com/TuringLang/Turing.jl/issues/2151
-        return eltype(vi[spl])
+        return eltype(vi[:])
     end
     return eltype(T)
 end
@@ -720,25 +718,11 @@ end
 
 # Utilities
 """
-    unflatten(vi::AbstractVarInfo[, context::AbstractContext], x::AbstractVector)
+    unflatten(vi::AbstractVarInfo, x::AbstractVector)
 
 Return a new instance of `vi` with the values of `x` assigned to the variables.
-
-If `context` is provided, `x` is assumed to be realizations only for variables not
-filtered out by `context`.
 """
-function unflatten(varinfo::AbstractVarInfo, context::AbstractContext, θ)
-    if hassampler(context)
-        unflatten(getsampler(context), varinfo, context, θ)
-    else
-        DynamicPPL.unflatten(varinfo, θ)
-    end
-end
-
-# TODO: deprecate this once `sampler` is no longer the main way of filtering out variables.
-function unflatten(sampler::AbstractSampler, varinfo::AbstractVarInfo, ::AbstractContext, θ)
-    return unflatten(varinfo, sampler, θ)
-end
+function unflatten end
 
 """
     to_maybe_linked_internal(vi::AbstractVarInfo, vn::VarName, dist, val)

src/compiler.jl

@@ -3,7 +3,7 @@ const INTERNALNAMES = (:__model__, :__context__, :__varinfo__)
 """
     need_concretize(expr)
 
-Return `true` if `expr` needs to be concretized, i.e., if it contains a colon `:` or 
+Return `true` if `expr` needs to be concretized, i.e., if it contains a colon `:` or
 requires a dynamic optic.
 
 # Examples
@@ -730,19 +730,19 @@ function warn_empty(body)
     return nothing
 end
 
+# TODO(mhauru) matchingvalue has methods that can accept both types and values. Why?
+# TODO(mhauru) This function needs a more comprehensive docstring.
 """
-    matchingvalue(sampler, vi, value)
-    matchingvalue(context::AbstractContext, vi, value)
+    matchingvalue(vi, value)
 
-Convert the `value` to the correct type for the `sampler` or `context` and the `vi` object.
-
-For a `context` that is _not_ a `SamplingContext`, we fall back to
-`matchingvalue(SampleFromPrior(), vi, value)`.
+Convert the `value` to the correct type for the `vi` object.
 """
-function matchingvalue(sampler, vi, value)
+function matchingvalue(vi, value)
     T = typeof(value)
     if hasmissing(T)
-        _value = convert(get_matching_type(sampler, vi, T), value)
+        _value = convert(get_matching_type(vi, T), value)
+        # TODO(mhauru) Why do we make a deepcopy, even though in the !hasmissing branch we
+        # are happy to return `value` as-is?
         if _value === value
             return deepcopy(_value)
         else
@@ -752,45 +752,30 @@ function matchingvalue(sampler, vi, value)
         return value
     end
 end
-# If we hit `Type` or `TypeWrap`, we immediately jump to `get_matching_type`.
-function matchingvalue(sampler::AbstractSampler, vi, value::FloatOrArrayType)
-    return get_matching_type(sampler, vi, value)
-end
-function matchingvalue(sampler::AbstractSampler, vi, value::TypeWrap{T}) where {T}
-    return TypeWrap{get_matching_type(sampler, vi, T)}()
-end
 
-function matchingvalue(context::AbstractContext, vi, value)
-    return matchingvalue(NodeTrait(matchingvalue, context), context, vi, value)
+function matchingvalue(vi, value::FloatOrArrayType)
+    return get_matching_type(vi, value)
 end
-function matchingvalue(::IsLeaf, context::AbstractContext, vi, value)
-    return matchingvalue(SampleFromPrior(), vi, value)
-end
-function matchingvalue(::IsParent, context::AbstractContext, vi, value)
-    return matchingvalue(childcontext(context), vi, value)
-end
-function matchingvalue(context::SamplingContext, vi, value)
-    return matchingvalue(context.sampler, vi, value)
+function matchingvalue(vi, ::TypeWrap{T}) where {T}
+    return TypeWrap{get_matching_type(vi, T)}()
 end
 
+# TODO(mhauru) This function needs a more comprehensive docstring. What is it for?
 """
-    get_matching_type(spl::AbstractSampler, vi, ::TypeWrap{T}) where {T}
-
-Get the specialized version of type `T` for sampler `spl`.
+    get_matching_type(vi, ::TypeWrap{T}) where {T}
 
-For example, if `T === Float64` and `spl::Hamiltonian`, the matching type is
-`eltype(vi[spl])`.
+Get the specialized version of type `T` for `vi`.
 """
-get_matching_type(spl::AbstractSampler, vi, ::Type{T}) where {T} = T
-function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Union{Missing,AbstractFloat}})
-    return Union{Missing,float_type_with_fallback(eltype(vi, spl))}
+get_matching_type(_, ::Type{T}) where {T} = T
+function get_matching_type(vi, ::Type{<:Union{Missing,AbstractFloat}})
+    return Union{Missing,float_type_with_fallback(eltype(vi))}
 end
-function get_matching_type(spl::AbstractSampler, vi, ::Type{<:AbstractFloat})
-    return float_type_with_fallback(eltype(vi, spl))
+function get_matching_type(vi, ::Type{<:AbstractFloat})
+    return float_type_with_fallback(eltype(vi))
 end
-function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T,N}}) where {T,N}
-    return Array{get_matching_type(spl, vi, T),N}
+function get_matching_type(vi, ::Type{<:Array{T,N}}) where {T,N}
+    return Array{get_matching_type(vi, T),N}
 end
-function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T}}) where {T}
-    return Array{get_matching_type(spl, vi, T)}
+function get_matching_type(vi, ::Type{<:Array{T}}) where {T}
+    return Array{get_matching_type(vi, T)}
 end

src/logdensityfunction.jl

@@ -121,22 +121,17 @@ end
 getsampler(f::LogDensityFunction) = getsampler(getcontext(f))
 hassampler(f::LogDensityFunction) = hassampler(getcontext(f))
 
-_get_indexer(ctx::AbstractContext) = _get_indexer(NodeTrait(ctx), ctx)
-_get_indexer(ctx::SamplingContext) = ctx.sampler
-_get_indexer(::IsParent, ctx::AbstractContext) = _get_indexer(childcontext(ctx))
-_get_indexer(::IsLeaf, ctx::AbstractContext) = Colon()
-
 """
     getparams(f::LogDensityFunction)
 
 Return the parameters of the wrapped varinfo as a vector.
 """
-getparams(f::LogDensityFunction) = f.varinfo[_get_indexer(getcontext(f))]
+getparams(f::LogDensityFunction) = f.varinfo[:]
 
 # LogDensityProblems interface
 function LogDensityProblems.logdensity(f::LogDensityFunction, θ::AbstractVector)
     context = getcontext(f)
-    vi_new = unflatten(f.varinfo, context, θ)
+    vi_new = unflatten(f.varinfo, θ)
     return getlogp(last(evaluate!!(f.model, vi_new, context)))
 end
 function LogDensityProblems.capabilities(::Type{<:LogDensityFunction})

src/model.jl

@@ -948,9 +948,9 @@ Return the arguments and keyword arguments to be passed to the evaluator of the
 ) where {_F,argnames}
     unwrap_args = [
         if is_splat_symbol(var)
-            :($matchingvalue(context_new, varinfo, model.args.$var)...)
+            :($matchingvalue(varinfo, model.args.$var)...)
         else
-            :($matchingvalue(context_new, varinfo, model.args.$var))
+            :($matchingvalue(varinfo, model.args.$var))
         end for var in argnames
     ]
 

src/sampler.jl

@@ -118,7 +118,7 @@ function AbstractMCMC.step(
 
     # Update the parameters if provided.
     if initial_params !== nothing
-        vi = initialize_parameters!!(vi, initial_params, spl, model)
+        vi = initialize_parameters!!(vi, initial_params, model)
 
         # Update joint log probability.
         # This is a quick fix for https://github.com/TuringLang/Turing.jl/issues/1588
@@ -156,9 +156,7 @@ By default, it returns an instance of [`SampleFromPrior`](@ref).
 """
 initialsampler(spl::Sampler) = SampleFromPrior()
 
-function set_values!!(
-    varinfo::AbstractVarInfo, initial_params::AbstractVector, spl::AbstractSampler
-)
+function set_values!!(varinfo::AbstractVarInfo, initial_params::AbstractVector)
     throw(
         ArgumentError(
             "`initial_params` must be a vector of type `Union{Real,Missing}`. " *
@@ -168,11 +166,9 @@ function set_values!!(
 end
 
 function set_values!!(
-    varinfo::AbstractVarInfo,
-    initial_params::AbstractVector{<:Union{Real,Missing}},
-    spl::AbstractSampler,
+    varinfo::AbstractVarInfo, initial_params::AbstractVector{<:Union{Real,Missing}}
 )
-    flattened_param_vals = varinfo[spl]
+    flattened_param_vals = varinfo[:]
     length(flattened_param_vals) == length(initial_params) || throw(
         DimensionMismatch(
             "Provided initial value size ($(length(initial_params))) doesn't match " *
@@ -189,12 +185,11 @@ function set_values!!(
     end
 
     # Update in `varinfo`.
-    return setindex!!(varinfo, flattened_param_vals, spl)
+    setall!(varinfo, flattened_param_vals)
+    return varinfo
 end
 
-function set_values!!(
-    varinfo::AbstractVarInfo, initial_params::NamedTuple, spl::AbstractSampler
-)
+function set_values!!(varinfo::AbstractVarInfo, initial_params::NamedTuple)
     vars_in_varinfo = keys(varinfo)
     for v in keys(initial_params)
         vn = VarName{v}()
@@ -219,23 +214,21 @@ function set_values!!(
     )
 end
 
-function initialize_parameters!!(
-    vi::AbstractVarInfo, initial_params, spl::AbstractSampler, model::Model
-)
+function initialize_parameters!!(vi::AbstractVarInfo, initial_params, model::Model)
     @debug "Using passed-in initial variable values" initial_params
 
     # `link` the varinfo if needed.
-    linked = islinked(vi, spl)
+    linked = islinked(vi)
     if linked
-        vi = invlink!!(vi, spl, model)
+        vi = invlink!!(vi, model)
     end
 
     # Set the values in `vi`.
-    vi = set_values!!(vi, initial_params, spl)
+    vi = set_values!!(vi, initial_params)
 
     # `invlink` if needed.
     if linked
-        vi = link!!(vi, spl, model)
+        vi = link!!(vi, model)
     end
 
     return vi

src/simple_varinfo.jl

@@ -258,7 +258,6 @@ function typed_simple_varinfo(model::Model)
     return last(evaluate!!(model, varinfo, SamplingContext()))
 end
 
-unflatten(svi::SimpleVarInfo, spl::AbstractSampler, x::AbstractVector) = unflatten(svi, x)
 function unflatten(svi::SimpleVarInfo, x::AbstractVector)
     logp = getlogp(svi)
     vals = unflatten(svi.values, x)
@@ -342,10 +341,6 @@ function BangBang.setindex!!(vi::SimpleVarInfo, val, vn::VarName)
     return Accessors.@set vi.values = set!!(vi.values, vn, val)
 end
 
-function BangBang.setindex!!(vi::SimpleVarInfo, val, spl::AbstractSampler)
-    return unflatten(vi, spl, val)
-end
-
 # TODO: Specialize to handle certain cases, e.g. a collection of `VarName` with
 # same symbol and same type of, say, `IndexLens`, for improved `.~` performance.
 function BangBang.setindex!!(vi::SimpleVarInfo, vals, vns::AbstractVector{<:VarName})
@@ -428,11 +423,7 @@ const SimpleOrThreadSafeSimple{T,V,C} = Union{
 }
 
 # Necessary for `matchingvalue` to work properly.
-function Base.eltype(
-    vi::SimpleOrThreadSafeSimple{<:Any,V}, spl::Union{AbstractSampler,SampleFromPrior}
-) where {V}
-    return V
-end
+Base.eltype(::SimpleOrThreadSafeSimple{<:Any,V}) where {V} = V
 
 # `subset`
 function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
@@ -562,7 +553,7 @@ istrans(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)
 istrans(vi::SimpleVarInfo, vn::VarName) = istrans(vi)
 istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, vn::VarName) = istrans(vi.varinfo, vn)
 
-islinked(vi::SimpleVarInfo, ::Union{Sampler,SampleFromPrior}) = istrans(vi)
+islinked(vi::SimpleVarInfo) = istrans(vi)
 
 values_as(vi::SimpleVarInfo) = vi.values
 values_as(vi::SimpleVarInfo{<:T}, ::Type{T}) where {T} = vi.values

src/threadsafe.jl

@@ -79,7 +79,7 @@ setval!(vi::ThreadSafeVarInfo, val, vn::VarName) = setval!(vi.varinfo, val, vn)
 keys(vi::ThreadSafeVarInfo) = keys(vi.varinfo)
 haskey(vi::ThreadSafeVarInfo, vn::VarName) = haskey(vi.varinfo, vn)
 
-islinked(vi::ThreadSafeVarInfo, spl::AbstractSampler) = islinked(vi.varinfo, spl)
+islinked(vi::ThreadSafeVarInfo) = islinked(vi.varinfo)
 
 function link!!(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
     return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, args...)
@@ -138,17 +138,6 @@ end
 function getindex(vi::ThreadSafeVarInfo, vns::AbstractVector{<:VarName}, dist::Distribution)
     return getindex(vi.varinfo, vns, dist)
 end
-getindex(vi::ThreadSafeVarInfo, spl::AbstractSampler) = getindex(vi.varinfo, spl)
-
-function BangBang.setindex!!(vi::ThreadSafeVarInfo, val, spl::AbstractSampler)
-    return Accessors.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, val, spl)
-end
-function BangBang.setindex!!(vi::ThreadSafeVarInfo, val, spl::SampleFromPrior)
-    return Accessors.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, val, spl)
-end
-function BangBang.setindex!!(vi::ThreadSafeVarInfo, val, spl::SampleFromUniform)
-    return Accessors.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, val, spl)
-end
 
 function BangBang.setindex!!(vi::ThreadSafeVarInfo, vals, vn::VarName)
     return Accessors.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, vals, vn)
@@ -184,13 +173,9 @@ function is_flagged(vi::ThreadSafeVarInfo, vn::VarName, flag::String)
     return is_flagged(vi.varinfo, vn, flag)
 end
 
-# Transformations.
 function settrans!!(vi::ThreadSafeVarInfo, trans::Bool, vn::VarName)
     return Accessors.@set vi.varinfo = settrans!!(vi.varinfo, trans, vn)
 end
-function settrans!!(vi::ThreadSafeVarInfo, spl::AbstractSampler, dist::Distribution)
-    return Accessors.@set vi.varinfo = settrans!!(vi.varinfo, spl, dist)
-end
 
 istrans(vi::ThreadSafeVarInfo, vn::VarName) = istrans(vi.varinfo, vn)
 istrans(vi::ThreadSafeVarInfo, vns::AbstractVector{<:VarName}) = istrans(vi.varinfo, vns)
@@ -200,9 +185,6 @@ getindex_internal(vi::ThreadSafeVarInfo, vn::VarName) = getindex_internal(vi.var
 function unflatten(vi::ThreadSafeVarInfo, x::AbstractVector)
     return Accessors.@set vi.varinfo = unflatten(vi.varinfo, x)
 end
-function unflatten(vi::ThreadSafeVarInfo, spl::AbstractSampler, x::AbstractVector)
-    return Accessors.@set vi.varinfo = unflatten(vi.varinfo, spl, x)
-end
 
 function subset(varinfo::ThreadSafeVarInfo, vns::AbstractVector{<:VarName})
     return Accessors.@set varinfo.varinfo = subset(varinfo.varinfo, vns)