DI over Optimisers test

Author: klamike

Project.toml

@@ -3,11 +3,14 @@ uuid = "8c219513-6c13-44e5-85c4-1da37bbbae23"
 authors = ["Klamkin", "Michael <michael@klamkin.com> and contributors"]
 version = "1.0.0-DEV"
 
-[compat]
-julia = "1.6.7"
-
 [extras]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "DifferentiationInterface", "ForwardDiff", "Optimisers", "Zygote", "Enzyme", "Mooncake"]

test/runtests.jl

@@ -1,6 +1,65 @@
 using L2ODC3
 using Test
+using Optimisers
 
-@testset "L2ODC3.jl" begin
-    # Write your tests here.
-end
+using ForwardDiff
+using Zygote
+using Enzyme
+using Mooncake
+
+using DifferentiationInterface
+const DI = DifferentiationInterface
+
+
+@testset "DI over Optimisers" begin
+    function test_nested(outer, inner)
+        function run_adam_steps(N, x0)
+            x = copy(x0)
+            rule = Optimisers.Adam(0.1)
+            state = Optimisers.setup(rule, x)
+            for _ in 1:N
+                grad = DI.gradient(x -> sum(abs2.(x)), inner(), x)
+                state, x = Optimisers.update(state, x, grad)
+            end
+            return sum(x)
+        end
+
+        x0 = [1.0, 2.0, 3.0]
+        N = 5
+        run_adam_steps(N, x0)
+
+        grad = DI.gradient(x -> run_adam_steps(N, x), outer(), x0)
+
+        @test length(grad) == length(x0)
+        @test all(!isnan, grad)
+        @test all(>(0), grad)
+
+        true
+    end
+    backends = [AutoForwardDiff, AutoZygote, AutoEnzyme, AutoMooncake]
+    working_combos = Dict(k => Set() for k in backends)
+    for _outer in backends
+        for _inner in backends
+            @testset "$_outer over $_inner" begin
+                if (
+                    (_outer in (AutoZygote, AutoEnzyme, AutoMooncake)) ||
+                    (_outer === AutoForwardDiff && _inner in (
+                        AutoEnzyme, AutoMooncake
+                    ))
+                )
+                    @test_throws "" test_nested(_outer, _inner)
+                else
+                    @test test_nested(_outer, _inner)
+                    push!(working_combos[_outer], _inner)
+                end
+            end
+        end
+    end
+    ds = "Working combinations:\n"
+    for _outer in keys(working_combos)
+        ds *= "$_outer over: "
+        ds *= join(collect(working_combos[_outer]), ", ")
+        ds *= "\n"
+    end
+    @info ds
+end