fix Zygote support

Author: mohamed82008

src/compat/ad.jl

@@ -7,18 +7,12 @@ ZygoteRules.@adjoint function push!(
     dist::Distribution,
     gidset::Set{Selector}
 )
-    return push!(vi, vn, r, dist, gidset), _ -> nothing
+    return push!(vi, vn, r, dist, gidset), _ -> ntuple(_ -> nothing, 5)
 end
 
-# Multithreaded evaluation is not compatible with Zygote.
-ZygoteRules.@adjoint function (model::Model)(
-    vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    ctx::AbstractContext
-)
-    function evaluate(vi, spl, ctx)
-        return evaluate_singlethreaded(model, vi, spl, ctx)
-    end
-    return ZygoteRules.pullback(evaluate, vi, spl, ctx)
+ZygoteRules.@adjoint function Threads.nthreads()
+    Threads.nthreads(), _ -> (nothing,)
+end
+ZygoteRules.@adjoint function Threads.threadid()
+    Threads.threadid(), _ -> (nothing,)
 end
-

src/threadsafe.jl

@@ -1,3 +1,30 @@
+#################
+# VectorOfLogps #
+#################
+
+struct VectorOfLogps{T1, T2 <: Vector{Base.RefValue{T1}}}
+    v::T2
+end
+VectorOfLogps(::Type{T}, n::Int) where {T} = VectorOfLogps(zero(T), n)
+function VectorOfLogps(val::T, n::Int) where {T}
+    v = [val for i in 1:Threads.nthreads()]
+    return VectorOfLogps(v)
+end
+VectorOfLogps(v::Vector) = VectorOfLogps(Ref.(v))
+Base.getindex(v::VectorOfLogps, i::Integer) = v.v[i][]
+function Base.setindex!(v::VectorOfLogps, val, i::Integer)
+    v.v[i][] = val
+    return v
+end
+Base.sum(v::VectorOfLogps) = sum(v -> v[], v.v)
+function Base.fill!(v::VectorOfLogps, val)
+    for i in 1:length(v.v)
+        v.v[i][] = val
+    end
+    return v
+end
+
+
 """
     ThreadSafeVarInfo
 
@@ -9,7 +36,7 @@ struct ThreadSafeVarInfo{V<:AbstractVarInfo,L} <: AbstractVarInfo
     logps::L
 end
 function ThreadSafeVarInfo(vi::AbstractVarInfo)
-    return ThreadSafeVarInfo(vi, [zero(getlogp(vi)) for _ in 1:Threads.nthreads()])
+    return ThreadSafeVarInfo(vi, VectorOfLogps(zero(getlogp(vi)), Threads.nthreads()))
 end
 ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi