Remove Zygote dependency and update documentation for ForwardDiff integration

.github/workflows/AD.yml

@@ -30,7 +30,6 @@ jobs:
           - Mooncake
           - Tracker
           - ReverseDiff
-          - Zygote
     steps:
       - uses: actions/checkout@v4
 

Project.toml

@@ -29,7 +29,6 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
@@ -40,7 +39,6 @@ BijectorsMooncakeExt = "Mooncake"
 BijectorsReverseDiffExt = "ReverseDiff"
 BijectorsReverseDiffChainRulesExt = ["ChainRules", "ReverseDiff"]
 BijectorsTrackerExt = "Tracker"
-BijectorsZygoteExt = "Zygote"
 
 [compat]
 ArgCheck = "1, 2"
@@ -64,7 +62,6 @@ ReverseDiff = "1"
 Roots = "1.3.15, 2"
 Statistics = "1"
 Tracker = "0.2"
-Zygote = "0.6.63, 0.7"
 julia = "1.10.8"
 
 [extras]
@@ -75,4 +72,3 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

docs/Project.toml

@@ -2,10 +2,8 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Documenter = "0.27"
 Functors = "0.3"
-StableRNGs = "1"
-Zygote = "0.6"
+StableRNGs = "1"

docs/src/examples.md

@@ -112,7 +112,7 @@ y = rand(rng, td)
 Want to fit the flow?
 
 ```@repl normalizing-flows
-using Zygote
+using ForwardDiff
 
 # Construct the flow.
 b = PlanarLayer(2)
@@ -145,7 +145,7 @@ f = NLLObjective(reconstruct, MvNormal(2, 1), xs);
 # Train using gradient descent.
 ε = 1e-3;
 for i in 1:100
-    (∇s,) = Zygote.gradient(f, θs)
+    ∇s = ForwardDiff.gradient(θ -> f(θ), θs)
     θs = fmap(θs, ∇s) do θ, ∇
         θ - ε .* ∇
     end

src/bijectors/pd.jl

@@ -1,6 +1,6 @@
 struct PDBijector <: Bijector end
 
-# This function has custom adjoints defined for Tracker, Zygote and ReverseDiff.
+# This function has custom adjoints defined for Tracker and ReverseDiff.
 # I couldn't find a mutation-free implementation that maintains TrackedArrays in Tracker
 # and ReverseDiff, hence the need for custom adjoints.
 function replace_diag(f, X)

src/chainrules.jl

@@ -286,5 +286,48 @@ function ChainRulesCore.rrule(::typeof(pd_from_upper), X::AbstractMatrix)
     end
 end
 
-# Fixes Zygote's issues with `@debug`
+# Fixes AD issues with `@debug`
 ChainRulesCore.@non_differentiable _debug(::Any)
+
+# ChainRules for utility functions used by PlanarLayer
+function ChainRulesCore.rrule(::typeof(aT_b), a::AbstractVector, b::AbstractMatrix)
+    y = aT_b(a, b)
+    function aT_b_pullback(Δy)
+        # For y = a' * b, we have:
+        # ∂y/∂a = b * Δy (since y is 1×n and a is m×1)
+        # ∂y/∂b = a * Δy (since y is 1×n and b is m×n)
+        Δa = b * vec(Δy)
+        Δb = a * Δy
+        return ChainRulesCore.NoTangent(), Δa, Δb
+    end
+    return y, aT_b_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(aT_b), a::AbstractVector, b::AbstractVector)
+    y = aT_b(a, b)
+    function aT_b_pullback(Δy)
+        # For y = dot(a, b), we have:
+        # ∂y/∂a = b * Δy
+        # ∂y/∂b = a * Δy
+        return ChainRulesCore.NoTangent(), b * Δy, a * Δy
+    end
+    return y, aT_b_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(_vec), x::AbstractArray)
+    y = _vec(x)
+    function _vec_pullback(Δy)
+        # Reshape the gradient back to the original shape
+        Δx = reshape(Δy, size(x))
+        return ChainRulesCore.NoTangent(), Δx
+    end
+    return y, _vec_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(_vec), x::Real)
+    y = _vec(x)
+    function _vec_pullback(Δy)
+        return ChainRulesCore.NoTangent(), Δy
+    end
+    return y, _vec_pullback
+end

test/Project.toml

@@ -25,7 +25,6 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 AbstractMCMC = "5"
@@ -50,5 +49,4 @@ Mooncake = "0.4"
 ReverseDiff = "1.4.2"
 StableRNGs = "1"
 Tracker = "0.2.11"
-Zygote = "0.6.63, 0.7"
 julia = "1.10"

test/ad/flows.jl

@@ -1,26 +1,26 @@
 @testset "PlanarLayer" begin
     # logpdf of a flow with a planar layer and two-dimensional inputs
-    test_ad(randn(7)) do θ
+    test_ad(randn(7), (:EnzymeForward,)) do θ
         layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
         flow = transformed(MvNormal(zeros(2), I), layer)
         x = θ[6:7]
         return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)
     end
-    test_ad(randn(11)) do θ
+    test_ad(randn(11), (:EnzymeForward,)) do θ
         layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
         flow = transformed(MvNormal(zeros(2), I), layer)
         x = reshape(θ[6:end], 2, :)
         return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x))
     end
 
     # logpdf of a flow with the inverse of a planar layer and two-dimensional inputs
-    test_ad(randn(7)) do θ
+    test_ad(randn(7), (:EnzymeForward,)) do θ
         layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
         flow = transformed(MvNormal(zeros(2), I), inverse(layer))
         x = θ[6:7]
         return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)
     end
-    test_ad(randn(11)) do θ
+    test_ad(randn(11), (:EnzymeForward,)) do θ
         layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
         flow = transformed(MvNormal(zeros(2), I), inverse(layer))
         x = reshape(θ[6:end], 2, :)

test/ad/utils.jl

@@ -6,7 +6,6 @@ function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6)
         if !(
             b in (
                 :ForwardDiff,
-                :Zygote,
                 :Mooncake,
                 :ReverseDiff,
                 :Enzyme,
@@ -32,16 +31,6 @@ function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6)
         end
     end
 
-    if AD == "All" || AD == "Zygote"
-        if :Zygote in broken
-            @test_broken Zygote.gradient(f, x)[1] ≈ finitediff rtol = rtol atol = atol
-        else
-            ∇zygote = Zygote.gradient(f, x)[1]
-            @test (all(iszero, finitediff) && ∇zygote === nothing) ||
-                isapprox(∇zygote, finitediff; rtol=rtol, atol=atol)
-        end
-    end
-
     if AD == "All" || AD == "ReverseDiff"
         if :ReverseDiff in broken
             @test_broken ReverseDiff.gradient(f, x) ≈ finitediff rtol = rtol atol = atol