First test scenarios for Flux gradients (#352)

* Flux-Zygote : creating functions for test scenario

* Weakdeps

* Shorter workflows

* Exclude

* Exclude more

* Working test on Zygote

* Fix

* Fix Flux tests

* Enzyme working

* Fixes

* No random

* Reactivate CI

* Printing

* Runtime activity

* Make mergeable

* Add FiniteDifferences

* Imports

* Group weird tests

* define DIT

* Higher tolerance in test

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: nialamarcotte

.github/workflows/Test.yml

@@ -48,6 +48,7 @@ jobs:
           - Back/Zygote
           - Down/Detector
           - Down/DifferentiateWith
+          - Down/Flux
         exclude:
           # lts
           - version: 'lts'
@@ -68,8 +69,10 @@ jobs:
             group: Back/Symbolics
           - version: 'lts'
             group: Back/Tapir
-          - version: 'pre'
+          - version: 'lts'
             group: Down/Detector
+          - version: 'lts'
+            group: Down/Flux
           # pre-release
           - version: 'pre'
             group: Formalities
@@ -81,6 +84,8 @@ jobs:
             group: Back/SecondOrder
           - version: 'pre'
             group: Down/Detector
+          - version: 'pre'
+            group: Down/Flux
 
     steps:
       - uses: actions/checkout@v4
@@ -124,10 +129,12 @@ jobs:
           - Formalities
           - Zero
           - ForwardDiff
-          - Zygote
+          - Weird
         exclude:
           - version: 'lts'
             group: Formalities
+          - version: 'lts'
+            group: Weird
 
     steps:
       - uses: actions/checkout@v4

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -35,7 +35,8 @@ using Enzyme:
     gradient,
     gradient!,
     jacobian,
-    make_zero
+    make_zero,
+    make_zero!
 
 struct AutoDeferredEnzyme{M} <: ADTypes.AbstractADType
     mode::M

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -26,7 +26,7 @@ function DI.value_and_pullback(
     f,
     backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}},
     x::Number,
-    dy::AbstractArray,
+    dy,
     ::NoPullbackExtras,
 )
     tf, tx = typeof(f), typeof(x)
@@ -40,11 +40,28 @@ end
 function DI.value_and_pullback(
     f,
     backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}},
-    x::AbstractArray,
-    dy,
-    extras::NoPullbackExtras,
+    x,
+    dy::Number,
+    ::NoPullbackExtras,
 )
-    dx = similar(x)
+    dx_sametype = make_zero(x)
+    x_and_dx = Duplicated(x, dx_sametype)
+    _, y = if backend isa AutoDeferredEnzyme
+        autodiff_deferred(ReverseWithPrimal, Const(f), Active, x_and_dx)
+    else
+        autodiff(ReverseWithPrimal, Const(f), Active, x_and_dx)
+    end
+    if !isone(dy)
+        # TODO: generalize beyond Arrays?
+        dx_sametype .*= dy
+    end
+    return y, dx_sametype
+end
+
+function DI.value_and_pullback(
+    f, backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}}, x, dy, extras::NoPullbackExtras
+)
+    dx = make_zero(x)
     return DI.value_and_pullback!(f, dx, backend, x, dy, extras)
 end
 
@@ -60,36 +77,34 @@ function DI.value_and_pullback!(
     f,
     dx,
     backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}},
-    x::AbstractArray,
+    x,
     dy::Number,
     ::NoPullbackExtras,
 )
     dx_sametype = convert(typeof(x), dx)
-    dx_sametype .= zero(eltype(x))
+    make_zero!(dx_sametype)
     x_and_dx = Duplicated(x, dx_sametype)
     _, y = if backend isa AutoDeferredEnzyme
         autodiff_deferred(ReverseWithPrimal, Const(f), Active, x_and_dx)
     else
         autodiff(ReverseWithPrimal, Const(f), Active, x_and_dx)
     end
-    dx_sametype .*= dy
+    if !isone(dy)
+        # TODO: generalize beyond Arrays?
+        dx_sametype .*= dy
+    end
     return y, copyto!(dx, dx_sametype)
 end
 
 function DI.value_and_pullback!(
-    f,
-    dx,
-    backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}},
-    x::AbstractArray,
-    dy::AbstractArray,
-    ::NoPullbackExtras,
+    f, dx, backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}}, x, dy, ::NoPullbackExtras
 )
     tf, tx = typeof(f), typeof(x)
     forw, rev = autodiff_thunk(
         ReverseSplitWithPrimal, Const{tf}, Duplicated, Duplicated{tx}
     )
     dx_sametype = convert(typeof(x), dx)
-    dx_sametype .= zero(eltype(x))
+    make_zero!(dx_sametype)
     tape, y, new_dy = forw(Const(f), Duplicated(x, dx_sametype))
     copyto!(new_dy, dy)
     rev(Const(f), Duplicated(x, dx_sametype), tape)
@@ -133,7 +148,7 @@ function DI.gradient!(
     extras::NoGradientExtras,
 )
     grad_sametype = convert(typeof(x), grad)
-    grad_sametype .= zero(eltype(x))
+    make_zero!(grad_sametype)
     if backend isa AutoDeferredEnzyme
         autodiff_deferred(reverse_mode(backend), f, Active, Duplicated(x, grad_sametype))
     else
@@ -145,13 +160,13 @@ end
 function DI.value_and_gradient(
     f, backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}}, x, ::NoGradientExtras
 )
-    return DI.value_and_pullback(f, backend, x, one(eltype(x)), NoPullbackExtras())
+    return DI.value_and_pullback(f, backend, x, true, NoPullbackExtras())
 end
 
 function DI.value_and_gradient!(
     f, grad, backend::AnyAutoEnzyme{<:Union{ReverseMode,Nothing}}, x, ::NoGradientExtras
 )
-    return DI.value_and_pullback!(f, grad, backend, x, one(eltype(x)), NoPullbackExtras())
+    return DI.value_and_pullback!(f, grad, backend, x, true, NoPullbackExtras())
 end
 
 ## Jacobian

DifferentiationInterface/test/Down/Flux/Project.toml

@@ -0,0 +1,7 @@
+[deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

DifferentiationInterface/test/Down/Flux/test.jl

@@ -0,0 +1,21 @@
+using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterfaceTest as DIT
+using FiniteDifferences: FiniteDifferences
+using Flux: Flux
+using Enzyme: Enzyme
+using Zygote: Zygote
+using Test
+
+Enzyme.API.runtimeActivity!(true)
+
+test_differentiation(
+    [
+        AutoZygote(),
+        # AutoEnzyme()  # TODO: fix
+    ],
+    flux_scenarios();
+    isequal=DIT.flux_isequal,
+    isapprox=DIT.flux_isapprox,
+    rtol=1e-2,
+    atol=1e-6,
+)

DifferentiationInterface/test/runtests.jl

@@ -24,7 +24,7 @@ function subtest(category, folder)
     @testset "$file" for file in filter(
         endswith(".jl"), readdir(joinpath(@__DIR__, category, folder))
     )
-        @info "Testing category/$folder/$file"
+        @info "Testing $category/$folder/$file"
         include(joinpath(@__DIR__, category, folder, file))
     end
     Pkg.activate(TEST_ENV)

DifferentiationInterfaceTest/Project.toml

@@ -10,6 +10,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
@@ -20,11 +21,14 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 DifferentiationInterfaceTestComponentArraysExt = "ComponentArrays"
+DifferentiationInterfaceTestFluxExt = ["FiniteDifferences", "Flux"]
 DifferentiationInterfaceTestJLArraysExt = "JLArrays"
 DifferentiationInterfaceTestStaticArraysExt = "StaticArrays"
 
@@ -36,6 +40,7 @@ ComponentArrays = "0.15"
 DataFrames = "1.6.1"
 DifferentiationInterface = "0.5.6"
 DocStringExtensions = "0.9"
+Functors = "0.4"
 JET = "0.4 - 0.8, 0.9"
 JLArrays = "0.1"
 LinearAlgebra = "<0.0.1,1"
@@ -54,6 +59,8 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
@@ -73,6 +80,8 @@ test = [
     "ComponentArrays",
     "DataFrames",
     "DifferentiationInterface",
+    "FiniteDifferences",
+    "Flux",
     "ForwardDiff",
     "JET",
     "JLArrays",

DifferentiationInterfaceTest/docs/src/api.md

@@ -25,6 +25,7 @@ sparse_scenarios
 component_scenarios
 gpu_scenarios
 static_scenarios
+flux_scenarios
 ```
 
 ## Scenario types