Remove Zygote dependency and update documentation for ForwardDiff integration (#393)

* Remove Zygote dependency and update documentation for ForwardDiff integration

* Remove Zygote from AD workflow jobs

* fixed failing tests

* bump to remove the expected broken things|

* tests pls

* Add ChainRules for utility functions in PlanarLayer and update tests for EnzymeForward

* remove things to try and pass the thing

---------

Co-authored-by: AoifeHughes <[email protected]>

Author: AoifeHughes

.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

ext/BijectorsZygoteExt.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/bijectors/pd.jl

@@ -286,5 +286,6 @@ 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)
+

src/chainrules.jl

@@ -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/Project.toml

@@ -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

test/ad/utils.jl

@@ -13,7 +13,6 @@ using LogExpFunctions
 using Mooncake
 using ReverseDiff
 using Tracker
-using Zygote
 
 using Random, LinearAlgebra, Test