Make ZygoteRules and ChainRulesCore weak dependencies

Project.toml

@@ -23,10 +23,14 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [extensions]
+DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
+DynamicPPLZygoteRulesExt = ["ZygoteRules"]
 
 [compat]
 AbstractMCMC = "5"
@@ -51,4 +55,6 @@ Test = "1.6"
 julia = "1.6"
 
 [extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

ext/DynamicPPLChainRulesCoreExt.jl

@@ -0,0 +1,27 @@
+module DynamicPPLChainRulesCoreExt
+
+if isdefined(Base, :get_extension)
+    using DynamicPPL: DynamicPPL, BangBang, Distributions
+    using ChainRulesCore: ChainRulesCore
+else
+    using ..DynamicPPL: DynamicPPL, BangBang, Distributions
+    using ..ChainRulesCore: ChainRulesCore
+end
+
+# See https://github.com/TuringLang/Turing.jl/issues/1199
+ChainRulesCore.@non_differentiable BangBang.push!!(
+    vi::DynamicPPL.VarInfo,
+    vn::DynamicPPL.VarName,
+    r,
+    dist::Distributions.Distribution,
+    gidset::Set{DynamicPPL.Selector},
+)
+
+ChainRulesCore.@non_differentiable DynamicPPL.updategid!(
+    vi::DynamicPPL.AbstractVarInfo, vn::DynamicPPL.VarName, spl::DynamicPPL.Sampler
+)
+
+# No need + causes issues for some AD backends, e.g. Zygote.
+ChainRulesCore.@non_differentiable DynamicPPL.infer_nested_eltype(x)
+
+end # module

ext/DynamicPPLZygoteRulesExt.jl

@@ -0,0 +1,25 @@
+module DynamicPPLZygoteRulesExt
+
+if isdefined(Base, :get_extension)
+    using DynamicPPL: DynamicPPL, Distributions
+    using ZygoteRules: ZygoteRules
+else
+    using ..DynamicPPL: DynamicPPL, Distributions
+    using ..ZygoteRules: ZygoteRules
+end
+
+# https://github.com/TuringLang/Turing.jl/issues/1595
+ZygoteRules.@adjoint function DynamicPPL.dot_observe(
+    spl::Union{DynamicPPL.SampleFromPrior,DynamicPPL.SampleFromUniform},
+    dists::AbstractArray{<:Distributions.Distribution},
+    value::AbstractArray,
+    vi,
+)
+    function dot_observe_fallback(spl, dists, value, vi)
+        DynamicPPL.increment_num_produce!(vi)
+        return sum(map(Distributions.loglikelihood, dists, value)), vi
+    end
+    return ZygoteRules.pullback(__context__, dot_observe_fallback, spl, dists, value, vi)
+end
+
+end # module

src/DynamicPPL.jl

@@ -9,11 +9,9 @@
 
 using AbstractMCMC: AbstractMCMC
 using BangBang: BangBang, push!!, empty!!, setindex!!
-using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ConstructionBase: ConstructionBase
 using Setfield: Setfield
-using ZygoteRules: ZygoteRules
 using LogDensityProblems: LogDensityProblems
 
 using LinearAlgebra: LinearAlgebra, Cholesky
@@ -171,7 +169,6 @@
 include("context_implementations.jl")
 include("compiler.jl")
 include("prob_macro.jl")
-include("compat/ad.jl")
 include("loglikelihoods.jl")
 include("submodel_macro.jl")
 include("test_utils.jl")
@@ -186,9 +183,15 @@
 
 @static if !isdefined(Base, :get_extension)
     function __init__()
+        @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" include(
+            "../ext/DynamicPPLChainRulesCoreExt.jl"
+        )
         @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(
             "../ext/DynamicPPLMCMCChainsExt.jl"
         )
+        @require ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" include(
+            "../ext/DynamicPPLZygoteRulesExt.jl"
+        )
     end
 end
 

src/utils.jl

@@ -883,9 +883,6 @@ 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)
-
 """
     varname_leaves(vn::VarName, val)