Use ChainRulesCore and define adjoint for updategid! (#163)

This PR moves the definition of the adjoint for `updategid!` to DynamicPPL and replaces ZygoteRules with ChainRulesCore, as mentioned in a comment in TuringLang/Turing.jl#1401.

Author: devmotion

Project.toml

@@ -1,21 +1,21 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.9.0"
+version = "0.9.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaturalSort = "c020b1a1-e9b0-503a-9c33-f039bfc54a85"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 AbstractMCMC = "1"
 Bijectors = "0.5.2, 0.6, 0.7, 0.8"
+ChainRulesCore = "0.9.7"
 Distributions = "0.23.8"
 MacroTools = "0.5.1"
 NaturalSort = "1"
-ZygoteRules = "0.2"
 julia = "1.3"

src/DynamicPPL.jl

@@ -5,9 +5,9 @@ using Distributions
 using Bijectors
 
 import AbstractMCMC
+import ChainRulesCore
 import NaturalSort
 import MacroTools
-import ZygoteRules
 
 import Random
 

src/compat/ad.jl

@@ -1,11 +1,14 @@
-# Prevent Zygote from differentiating push!
 # See https://github.com/TuringLang/Turing.jl/issues/1199
-ZygoteRules.@adjoint function push!(
+ChainRulesCore.@non_differentiable push!(
     vi::VarInfo,
     vn::VarName,
     r,
     dist::Distribution,
     gidset::Set{Selector}
 )
-    return push!(vi, vn, r, dist, gidset), _ -> nothing
-end
+
+ChainRulesCore.@non_differentiable updategid!(
+    vi::AbstractVarInfo,
+    vn::VarName,
+    spl::Sampler,
+)

test/Turing/core/compat/zygote.jl

@@ -26,5 +26,3 @@ function gradient_logp(
 
     return l, ∂l∂θ
 end
-
-Zygote.@nograd DynamicPPL.updategid!