fix AutoZygote gradient (#2664)

Author: CarloLucibello

src/gradient.jl

@@ -54,6 +54,8 @@ function gradient(f, adtype::AbstractADType, args...)
         Supported backends are $SUPPORTED_AD_BACKENDS.")
 end
 
+gradient(f, adtype::AutoZygote, args...) = Zygote.gradient(f, args...)
+
 # Default gradient using Zygote
 function gradient(f, args...; zero::Bool=true)
     for a in args

test/gradient.jl

@@ -0,0 +1,7 @@
+@testset "AutoADTypes gradient" begin
+    m = Dense(2 => 2)
+    x = rand(Float32, 2)
+    g_zygote = Flux.gradient(m -> sum(m(x)), AutoZygote(), m)[1]
+    g_mooncake = Flux.gradient(m -> sum(m(x)), AutoMooncake(), m)[1]
+    @test g_zygote.weight ≈ g_mooncake.weight
+end