Fix Zygote issue with dot_observe (#236)

This PR fixes TuringLang/Turing.jl#1595.

It is an alternative to #235 that does not require us to rewrite the primal less efficiently which would affect regular execution and other AD backends.

Co-authored-by: David Widmann <[email protected]>

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.10.16"
+version = "0.10.17"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -11,6 +11,7 @@ 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 = "2, 3.0"
@@ -20,4 +21,5 @@ ChainRulesCore = "0.9.7"
 Distributions = "0.23.8, 0.24"
 MacroTools = "0.5.6"
 NaturalSort = "1"
+ZygoteRules = "0.2"
 julia = "1.3"

src/DynamicPPL.jl

@@ -9,6 +9,7 @@ import AbstractMCMC
 import ChainRulesCore
 import NaturalSort
 import MacroTools
+import ZygoteRules
 
 import Random
 

src/compat/ad.jl

@@ -12,3 +12,17 @@ ChainRulesCore.@non_differentiable updategid!(
     vn::VarName,
     spl::Sampler,
 )
+
+# https://github.com/TuringLang/Turing.jl/issues/1595
+ZygoteRules.@adjoint function dot_observe(
+    spl::Union{SampleFromPrior, SampleFromUniform},
+    dists::AbstractArray{<:Distribution},
+    value::AbstractArray,
+    vi,
+)
+    function dot_observe_fallback(spl, dists, value, vi)
+        increment_num_produce!(vi)
+        return sum(map(Distributions.loglikelihood, dists, value))
+    end
+    return ZygoteRules.pullback(__context__, dot_observe_fallback, spl, dists, value, vi)
+end

test/compat/ad.jl

@@ -24,4 +24,29 @@
 
         test_model_ad(wishart_ad(), logp_wishart_ad)
     end
+
+    # https://github.com/TuringLang/Turing.jl/issues/1595
+    @testset "dot_observe" begin
+        function f_dot_observe(x)
+            return DynamicPPL.dot_observe(SampleFromPrior(), [Normal(), Normal(-1.0, 2.0)], x, VarInfo())
+        end
+        function f_dot_observe_manual(x)
+            return logpdf(Normal(), x[1]) + logpdf(Normal(-1.0, 2.0), x[2])
+        end
+
+        # Manual computation of the gradient.
+        x = randn(2)
+        val = f_dot_observe_manual(x)
+        grad = ForwardDiff.gradient(f_dot_observe_manual, x)
+
+        @test ForwardDiff.gradient(f_dot_observe, x) ≈ grad
+
+        y, back = Tracker.forward(f_dot_observe, x)
+        @test Tracker.data(y) ≈ val
+        @test Tracker.data(back(1)[1]) ≈ grad
+
+        y, back = Zygote.pullback(f_dot_observe, x)
+        @test y ≈ val
+        @test back(1)[1] ≈ grad
+    end
 end