Update ChainRulesCore (#172)

Author: devmotion

Project.toml

@@ -23,8 +23,8 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 Adapt = "2, 3"
-ChainRules = "0.7"
-ChainRulesCore = "0.9.21"
+ChainRules = "0.7, 0.8"
+ChainRulesCore = "0.9.44, 0.10"
 Compat = "3.6"
 DiffRules = "0.1, 1.0"
 Distributions = "0.23.3, 0.24, 0.25"

src/common.jl

@@ -7,7 +7,7 @@ end
 function turing_chol_back(A::AbstractMatrix, check)
     C, chol_pullback = rrule(cholesky, A, Val(false), check=check)
     function back(Δ)
-        Ȳ = Composite{typeof(C)}((U=Δ[1]))
+        Ȳ = Tangent{typeof(C)}((U=Δ[1]))
         ∂C = chol_pullback(Ȳ)[2]
         (∂C, nothing)
     end
@@ -21,7 +21,7 @@ end
 function symm_turing_chol_back(A::AbstractMatrix, check, uplo)
     C, chol_pullback = rrule(cholesky, Symmetric(A,uplo), Val(false), check=check)
     function back(Δ)
-        Ȳ = Composite{typeof(C)}((U=Δ[1]))
+        Ȳ = Tangent{typeof(C)}((U=Δ[1]))
         ∂C = chol_pullback(Ȳ)[2]
         (∂C, nothing, nothing)
     end

test/Project.toml

@@ -17,8 +17,8 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ChainRulesCore = "0.9"
-ChainRulesTestUtils = "0.6.3"
+ChainRulesCore = "0.9.44, 0.10"
+ChainRulesTestUtils = "0.6.3, 0.7"
 Combinatorics = "1.0.2"
 Distributions = "0.24.3, 0.25"
 FiniteDifferences = "0.11.3, 0.12"

test/ad/chainrules.jl

@@ -34,14 +34,15 @@
     test_frule(StatsFuns.tdistlogpdf, x, y)
     test_rrule(StatsFuns.tdistlogpdf, x, y)
 
+    # TODO: Re-enable if https://github.com/JuliaMath/SpecialFunctions.jl/pull/325 is fixed
     # use `BigFloat` to avoid Rmath implementation in finite differencing check
     # (returns `NaN` for non-integer values)
-    n = rand(1:100)
-    x = BigFloat(n)
-    y = big(logistic(randn()))
-    z = BigFloat(rand(1:n))
-    test_frule(StatsFuns.binomlogpdf, x, y, z)
-    test_rrule(StatsFuns.binomlogpdf, x, y, z)
+    #n = rand(1:100)
+    #x = BigFloat(n)
+    #y = big(logistic(randn()))
+    #z = BigFloat(rand(1:n))
+    #test_frule(StatsFuns.binomlogpdf, x, y, z)
+    #test_rrule(StatsFuns.binomlogpdf, x, y, z)
 
     x = big(exp(randn()))
     y = BigFloat(rand(1:100))

test/ad/distributions.jl

@@ -54,8 +54,8 @@
         DistSpec(Poisson, (0.5,), 1),
         DistSpec(Poisson, (0.5,), [1, 1]),
 
-        DistSpec(Skellam, (1.0, 2.0), -2; broken=(:Zygote,)),
-        DistSpec(Skellam, (1.0, 2.0), [-2, -2]; broken=(:Zygote,)),
+        DistSpec(Skellam, (1.0, 2.0), -2),
+        DistSpec(Skellam, (1.0, 2.0), [-2, -2]),
 
         DistSpec(PoissonBinomial, ([0.5, 0.5],), 0),
 
@@ -162,8 +162,9 @@
 
         DistSpec(NormalCanon, (1.0, 2.0), 0.5),
 
-        DistSpec(NormalInverseGaussian, (1.0, 2.0, 1.0, 1.0), 0.5; broken=(:Zygote,)),
+        DistSpec(NormalInverseGaussian, (1.0, 2.0, 1.0, 1.0), 0.5),
 
+        DistSpec(Pareto, (), 1.5),
         DistSpec(Pareto, (1.0,), 1.5),
         DistSpec(Pareto, (1.0, 1.0), 1.5),
 
@@ -213,10 +214,7 @@
         # Stackoverflow caused by SpecialFunctions.besselix
         DistSpec(VonMises, (1.0,), 1.0),
         DistSpec(VonMises, (1, 1), 1),
-        
-        # Only some Zygote tests are broken and therefore this can not be checked
-        DistSpec(Pareto, (), 1.5; broken=(:Zygote,)),
-        
+
         # Some tests are broken on some Julia versions, therefore it can't be checked reliably
         DistSpec(PoissonBinomial, ([0.5, 0.5],), [0, 0]; broken=(:Zygote,)),
     ]
@@ -395,17 +393,21 @@
             # Skellam only fails in these tests with ReverseDiff
             # Ref: https://github.com/TuringLang/DistributionsAD.jl/issues/126
             # PoissonBinomial fails with Zygote
+            # Matrix case does not work with Skellam:
+            # https://github.com/TuringLang/DistributionsAD.jl/pull/172#issuecomment-853721493
             filldist_broken = if d.f(d.θ...) isa Skellam
-                (d.broken..., :ReverseDiff)
+                ((d.broken..., :ReverseDiff), (d.broken..., :Zygote, :ReverseDiff))
             elseif d.f(d.θ...) isa PoissonBinomial
-                (d.broken..., :Zygote)
+                ((d.broken..., :Zygote), (d.broken..., :Zygote))
             else
-                d.broken
+                (d.broken, d.broken)
             end
-            arraydist_broken = if d.f(d.θ...) isa PoissonBinomial
-                (d.broken..., :Zygote)
+            arraydist_broken = if d.f(d.θ...) isa Skellam
+                (d.broken, (d.broken..., :Zygote))
+            elseif d.f(d.θ...) isa PoissonBinomial
+                ((d.broken..., :Zygote), (d.broken..., :Zygote))
             else
-                d.broken
+                (d.broken, d.broken)
             end
 
             # Create `filldist` distribution
@@ -416,7 +418,7 @@
             f_arraydist = (θ...,) -> arraydist([d.f(θ...) for _ in 1:n])
             d_arraydist = f_arraydist(d.θ...)
 
-            for sz in ((n,), (n, 2))
+            for (i, sz) in enumerate(((n,), (n, 2)))
                 # Matrix case doesn't work for continuous distributions for some reason
                 # now but not too important (?!)
                 if length(sz) == 2 && Distributions.value_support(typeof(d)) === Continuous
@@ -434,7 +436,7 @@
                         d.θ,
                         x,
                         d.xtrans;
-                        broken=filldist_broken,
+                        broken=filldist_broken[i],
                     )
                 )
                 test_ad(
@@ -444,7 +446,7 @@
                         d.θ,
                         x,
                         d.xtrans;
-                        broken=arraydist_broken,
+                        broken=arraydist_broken[i],
                     )
                 )
             end