Improvements to ThreadSafeVarInfo (#425)

The `ThreadSafeVarInfo` still has quite a few implementation details inherited from compat with `VarInfo`. This PR makes `ThreadSafeVarInfo` work better with other implementations, e.g. `SimpleVarInfo`.

It also fixes TuringLang/Turing.jl#1878 (comment)

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.20.1"
+version = "0.20.2"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/compiler.jl

@@ -684,17 +684,14 @@ For example, if `T === Float64` and `spl::Hamiltonian`, the matching type is
 """
 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,floatof(eltype(vi, spl))}
+    return Union{Missing,float_type_with_fallback(eltype(vi, spl))}
 end
 function get_matching_type(spl::AbstractSampler, vi, ::Type{<:AbstractFloat})
-    return floatof(eltype(vi, spl))
+    return float_type_with_fallback(eltype(vi, spl))
 end
 function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T,N}}) where {T,N}
     return Array{get_matching_type(spl, vi, T),N}
 end
 function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T}}) where {T}
     return Array{get_matching_type(spl, vi, T)}
 end
-
-floatof(::Type{T}) where {T<:Real} = typeof(one(T) / one(T))
-floatof(::Type) = Real # fallback if type inference failed

src/simple_varinfo.jl

@@ -202,18 +202,37 @@ struct SimpleVarInfo{NT,T,C<:AbstractTransformation} <: AbstractVarInfo
     transformation::C
 end
 
-SimpleVarInfo(values, logp) = SimpleVarInfo(values, logp, NoTransformation())
+# Makes things a bit more readable vs. putting `Float64` everywhere.
+const SIMPLEVARINFO_DEFAULT_ELTYPE = Float64
 
+function SimpleVarInfo{NT,T}(values, logp) where {NT,T}
+    return SimpleVarInfo{NT,T,NoTransformation}(values, logp, NoTransformation())
+end
 function SimpleVarInfo{T}(θ) where {T<:Real}
-    return SimpleVarInfo(θ, zero(T))
+    return SimpleVarInfo{typeof(θ),T}(θ, zero(T))
+end
+
+# Constructors without type-specification.
+SimpleVarInfo(θ) = SimpleVarInfo{SIMPLEVARINFO_DEFAULT_ELTYPE}(θ)
+function SimpleVarInfo(θ::Union{<:NamedTuple,<:AbstractDict})
+    return if isempty(θ)
+        # Can't infer from values, so we just use default.
+        SimpleVarInfo{SIMPLEVARINFO_DEFAULT_ELTYPE}(θ)
+    else
+        # Infer from `values`.
+        SimpleVarInfo{float_type_with_fallback(infer_nested_eltype(typeof(θ)))}(θ)
+    end
 end
+
+SimpleVarInfo(values, logp) = SimpleVarInfo{typeof(values),typeof(logp)}(values, logp)
+
+# Using `kwargs` to specify the values.
 function SimpleVarInfo{T}(; kwargs...) where {T<:Real}
     return SimpleVarInfo{T}(NamedTuple(kwargs))
 end
 function SimpleVarInfo(; kwargs...)
-    return SimpleVarInfo{Float64}(NamedTuple(kwargs))
+    return SimpleVarInfo(NamedTuple(kwargs))
 end
-SimpleVarInfo(θ) = SimpleVarInfo{Float64}(θ)
 
 # Constructor from `Model`.
 SimpleVarInfo(model::Model, args...) = SimpleVarInfo{Float64}(model, args...)
@@ -582,3 +601,12 @@ julia> # Truth.
 ```
 """
 Distributions.loglikelihood(model::Model, θ) = loglikelihood(model, SimpleVarInfo(θ))
+
+# Threadsafe stuff.
+# For `SimpleVarInfo` we don't really need `Ref` so let's not use it.
+function ThreadSafeVarInfo(vi::SimpleVarInfo)
+    return ThreadSafeVarInfo(vi, zeros(typeof(getlogp(vi)), Threads.nthreads()))
+end
+function ThreadSafeVarInfo(vi::SimpleVarInfo{<:Any,<:Ref})
+    return ThreadSafeVarInfo(vi, [Ref(zero(getlogp(vi))) for _ in 1:Threads.nthreads()])
+end

src/threadsafe.jl

@@ -13,26 +13,40 @@ function ThreadSafeVarInfo(vi::AbstractVarInfo)
 end
 ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi
 
+const ThreadSafeVarInfoWithRef{V<:AbstractVarInfo} = ThreadSafeVarInfo{
+    V,<:AbstractArray{<:Ref}
+}
+
 # Instead of updating the log probability of the underlying variables we
 # just update the array of log probabilities.
 function acclogp!!(vi::ThreadSafeVarInfo, logp)
+    vi.logps[Threads.threadid()] += logp
+    return vi
+end
+function acclogp!!(vi::ThreadSafeVarInfoWithRef, logp)
     vi.logps[Threads.threadid()][] += logp
     return vi
 end
 
 # The current log probability of the variables has to be computed from
 # both the wrapped variables and the thread-specific log probabilities.
-getlogp(vi::ThreadSafeVarInfo) = getlogp(vi.varinfo) + sum(getindex, vi.logps)
+getlogp(vi::ThreadSafeVarInfo) = getlogp(vi.varinfo) + sum(vi.logps)
+getlogp(vi::ThreadSafeVarInfoWithRef) = getlogp(vi.varinfo) + sum(getindex, vi.logps)
 
 # TODO: Make remaining methods thread-safe.
-
 function resetlogp!!(vi::ThreadSafeVarInfo)
+    return ThreadSafeVarInfo(resetlogp!!(vi.varinfo), zero(vi.logps))
+end
+function resetlogp!!(vi::ThreadSafeVarInfoWithRef)
     for x in vi.logps
         x[] = zero(x[])
     end
     return ThreadSafeVarInfo(resetlogp!!(vi.varinfo), vi.logps)
 end
 function setlogp!!(vi::ThreadSafeVarInfo, logp)
+    return ThreadSafeVarInfo(setlogp!!(vi.varinfo, logp), zero(vi.logps))
+end
+function setlogp!!(vi::ThreadSafeVarInfoWithRef, logp)
     for x in vi.logps
         x[] = zero(x[])
     end
@@ -104,7 +118,7 @@ end
 
 isempty(vi::ThreadSafeVarInfo) = isempty(vi.varinfo)
 function BangBang.empty!!(vi::ThreadSafeVarInfo)
-    return resetlogp!(Setfield.@set!(vi.varinfo = empty!!(vi.varinfo)))
+    return resetlogp!!(Setfield.@set!(vi.varinfo = empty!!(vi.varinfo)))
 end
 
 function BangBang.push!!(

src/utils.jl

@@ -546,3 +546,79 @@ function nested_haskey(dict::AbstractDict, vn::VarName)
 
     return canview(child, value)
 end
+
+"""
+    float_type_with_fallback(x)
+
+Return type corresponding to `float(typeof(x))` if possible; otherwise return `Real`.
+"""
+float_type_with_fallback(::Type) = Real
+float_type_with_fallback(::Type{T}) where {T<:Real} = float(T)
+
+"""
+    infer_nested_eltype(x::Type)
+
+Recursively unwrap the type, returning the first type where `eltype(x) === typeof(x)`.
+
+This is useful for obtaining a reasonable default `eltype` in deeply nested types.
+
+# Examples
+```jldoctest
+julia> # `AbstractArrary`
+       DynamicPPL.infer_nested_eltype(typeof([1.0]))
+Float64
+
+julia> # `NamedTuple` with `Float32`
+       DynamicPPL.infer_nested_eltype(typeof((x = [1f0], )))
+Float32
+
+julia> # `AbstractDict`
+       DynamicPPL.infer_nested_eltype(typeof(Dict(:x => [1.0, ])))
+Float64
+
+julia> # Nesting of containers.
+       DynamicPPL.infer_nested_eltype(typeof([Dict(:x => 1.0,) ]))
+Float64
+
+julia> DynamicPPL.infer_nested_eltype(typeof([Dict(:x => [1.0,],) ]))
+Float64
+
+julia> # Empty `Tuple`.
+       DynamicPPL.infer_nested_eltype(typeof(()))
+Any
+
+julia> # Empty `Dict`.
+       DynamicPPL.infer_nested_eltype(typeof(Dict()))
+Any
+```
+"""
+function infer_nested_eltype(::Type{T}) where {T}
+    ET = eltype(T)
+    return ET === T ? T : infer_nested_eltype(ET)
+end
+
+# We can do a better job than just `Any` with `Union`.
+infer_nested_eltype(::Type{Union{}}) = Any
+function infer_nested_eltype(::Type{U}) where {U<:Union}
+    return promote_type(U.a, infer_nested_eltype(U.b))
+end
+
+# Handle `NamedTuple` and `Tuple` specially given how prolific they are.
+function infer_nested_eltype(::Type{<:NamedTuple{<:Any,V}}) where {V}
+    return infer_nested_eltype(V)
+end
+
+# Recursively deal with `Tuple` so it has the potential of being compiled away.
+infer_nested_eltype(::Type{Tuple{T}}) where {T} = infer_nested_eltype(T)
+function infer_nested_eltype(::Type{T}) where {T<:Tuple{<:Any,Vararg{Any}}}
+    return promote_type(
+        infer_nested_eltype(Base.tuple_type_tail(T)),
+        infer_nested_eltype(Base.tuple_type_head(T)),
+    )
+end
+
+# Handle `AbstractDict` differently since `eltype` results in a `Pair`.
+infer_nested_eltype(::Type{<:AbstractDict{<:Any,ET}}) where {ET} = infer_nested_eltype(ET)
+
+# No need + causes issues for some AD backends, e.g. Zygote.
+ChainRulesCore.@non_differentiable infer_nested_eltype(x)