add static info for dist trace type

probcomp · georgematheos · May 17, 2020 · May 17, 2020 · May 17, 2020 · May 18, 2020
commit ffd9373c243593fbf46ef9b4331a93b9a787fdfd
diff --git a/src/distribution.jl b/src/distribution.jl
@@ -12,6 +12,9 @@ end
 abstract type Distribution{T} <: GenerativeFunction{T, DistributionTrace{T}} end
 @inline DistributionTrace(val::T, args::Tuple, dist::Dist) where {T, Dist <: Distribution} = DistributionTrace{T, Dist}(val, args, logpdf(dist, val, args...))
 
+# we need to know the specific distribution in the trace type so the compiler can specialize GFI calls fully
+@inline get_trace_type(::Dist) where {T, Dist <: Distribution{T}} = DistributionTrace{T, Dist}
+
 function Base.convert(::Type{DistributionTrace{U, <:Any}}, tr::DistributionTrace{<:Any, Dist}) where {U, Dist}
     DistributionTrace{U, Dist}(convert(U, tr.val), tr.args, dist(tr))
 end