Start replacing tests with DifferentiationInterface

.github/workflows/ci.yml

@@ -30,7 +30,7 @@ jobs:
         version:
           # - '1.9' # uncomment when julia 1.10 is out
           - '1'   # automatically expands to the latest stable 1.x release of Julia
-          - 'nightly'
+          # - 'nightly'
         os:
           - ubuntu-latest
           # - macOS-latest
@@ -40,19 +40,19 @@ jobs:
         julia-threads:
           - '1'
 
-        include:
-          - os: windows-latest
-            version: '1'
-            arch: x64
-            julia-threads: '1'
-          - os: macOS-latest
-            version: '1'
-            arch: x64
-            julia-threads: '1'
-          - os: ubuntu-latest
-            version: '1'
-            arch: x64
-            julia-threads: '2'
+        # include:
+        #   - os: windows-latest
+        #     version: '1'
+        #     arch: x64
+        #     julia-threads: '1'
+        #   - os: macOS-latest
+        #     version: '1'
+        #     arch: x64
+        #     julia-threads: '1'
+        #   - os: ubuntu-latest
+        #     version: '1'
+        #     arch: x64
+        #     julia-threads: '2'
 
     steps:
       - uses: actions/checkout@v3
@@ -79,29 +79,29 @@ jobs:
         with:
           file: lcov.info
 
-  docs:
-    name: Documentation
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v3
-      - uses: julia-actions/setup-julia@v1
-        with:
-          version: '1.9'
-      - uses: julia-actions/cache@v1
-      - run: |
-          julia --project=docs -e '
-            using Pkg
-            Pkg.develop(PackageSpec(path=pwd()))
-            Pkg.instantiate()'
-      - run: |
-          julia --color=yes --project=docs/ -e '
-            using NNlib
-            # using Pkg; Pkg.activate("docs")
-            using Documenter
-            using Documenter: doctest
-            DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib); recursive=true)
-            doctest(NNlib)'
-      - run: julia --project=docs docs/make.jl
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}
+  # docs:
+  #   name: Documentation
+  #   runs-on: ubuntu-latest
+  #   steps:
+  #     - uses: actions/checkout@v3
+  #     - uses: julia-actions/setup-julia@v1
+  #       with:
+  #         version: '1.9'
+  #     - uses: julia-actions/cache@v1
+  #     - run: |
+  #         julia --project=docs -e '
+  #           using Pkg
+  #           Pkg.develop(PackageSpec(path=pwd()))
+  #           Pkg.instantiate()'
+  #     - run: |
+  #         julia --color=yes --project=docs/ -e '
+  #           using NNlib
+  #           # using Pkg; Pkg.activate("docs")
+  #           using Documenter
+  #           using Documenter: doctest
+  #           DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib); recursive=true)
+  #           doctest(NNlib)'
+  #     - run: julia --project=docs docs/make.jl
+  #       env:
+  #         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  #         DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}

Project.toml

@@ -44,6 +44,7 @@ cuDNN = "1"
 julia = "1.9"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
@@ -64,4 +65,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [targets]
-test = ["AMDGPU", "CUDA", "ChainRulesTestUtils", "Documenter", "FiniteDifferences", "ForwardDiff", "Logging", "ReverseDiff", "StableRNGs", "Test", "UnicodePlots", "Zygote", "cuDNN", "Enzyme", "EnzymeCore", "EnzymeTestUtils", "Interpolations", "ImageTransformations"]
+test = ["ADTypes", "AMDGPU", "CUDA", "ChainRulesTestUtils", "Documenter", "FiniteDifferences", "ForwardDiff", "Logging", "ReverseDiff", "StableRNGs", "Test", "UnicodePlots", "Zygote", "cuDNN", "Enzyme", "EnzymeCore", "EnzymeTestUtils", "Interpolations", "ImageTransformations"]

test/runtests.jl

@@ -1,3 +1,12 @@
+using Pkg
+Pkg.add(url="https://github.com/gdalle/DifferentiationInterface.jl")
+Pkg.add(url="https://github.com/gdalle/DifferentiationInterface.jl", subdir="lib/DifferentiationInterfaceTest")
+
+import ADTypes
+import DifferentiationInterface as DI
+using DifferentiationInterfaceTest: Scenario
+import DifferentiationInterfaceTest as DIT
+
 using NNlib, Test, Statistics, Random
 using ChainRulesCore, ChainRulesTestUtils
 using Base.Broadcast: broadcasted
@@ -42,10 +51,12 @@ end
 
 cpu(x) = adapt(CPU(), x)
 
+#=
 include("gather.jl")
 include("scatter.jl")
 include("upsample.jl")
 include("rotation.jl")
+=#
 
 function nnlib_testsuite(Backend; skip_tests = Set{String}())
     @conditional_testset "Upsample" skip_tests begin
@@ -65,7 +76,8 @@ end
 @testset verbose=true "NNlib.jl" begin
 
     if get(ENV, "NNLIB_TEST_CPU", "true") == "true"
-        @testset "CPU" begin      
+        @testset verbose=true "CPU" begin      
+            #=
             @testset "Doctests" begin
                 doctest(NNlib, manual=false)
             end
@@ -113,11 +125,13 @@ end
             @testset "Padding" begin
                 include("padding.jl")
             end
+            =#
 
-            @testset "Softmax" begin
+            @testset verbose=true "Softmax" begin
                 include("softmax.jl")
             end
 
+            #=
             @testset "Utilities" begin
                 include("utils.jl")
             end
@@ -129,6 +143,7 @@ end
             @testset "Functions" begin
                 include("functions.jl")
             end
+            =#
         end
     else
         @info "Skipping CPU tests, set NNLIB_TEST_CPU=true to run them."

test/softmax.jl

@@ -110,21 +110,34 @@ end
     @test logsumexp(x; dims = 1) ≈ flogsoft(x, dims = 1)
 end
 
-@testset "AutoDiff" begin
+@testset verbose=true "AutoDiff" begin
+    scenarios = Scenario[]
+
     for f in (softmax, logsoftmax), d in (:, 1, 2)
-        gradtest(f, (3,4); fkwargs = (dims = d,), check_rrule = true)
+       push!(scenarios, Scenario(x -> f(x; dims=d), x=randn(3,4)))
     end
-    gradtest(x -> softmax(x) .* (1:3), 3)
-    gradtest(x -> softmax(x) .* (1:3), (3,5), atol = 1e-4)
-    gradtest(x -> softmax(x, dims = 2) .* (1:3), (3,5), atol = 1e-4)
+    push!(scenarios, Scenario(x -> softmax(x) .* (1:3), x=randn(3)))
+    push!(scenarios, Scenario(x -> softmax(x) .* (1:3), x=randn(3,5)))
+    push!(scenarios, Scenario(x -> softmax(x, dims = 2) .* (1:3), x=randn(3,5)))
 
-    gradtest(x -> logsoftmax(x) .* (1:3), 3)
-    gradtest(x -> logsoftmax(x) .* (1:3), (3,5))
-    gradtest(x -> logsoftmax(x, dims = 2) .* (1:3), (3,5))
+    push!(scenarios, Scenario(x -> logsoftmax(x) .* (1:3), x=randn(3)))
+    push!(scenarios, Scenario(x -> logsoftmax(x) .* (1:3), x=randn(3,5)))
+    push!(scenarios, Scenario(x -> logsoftmax(x, dims = 2) .* (1:3), x=randn(3,5)))
 
     for d  in (:, 1, 2)
-        gradtest(logsumexp, (3,4), fkwargs = (dims = d,))
+        push!(scenarios, Scenario(x -> logsumexp(x; dims=d), x=randn(3,4)))
     end
+
+    DIT.test_differentiation(
+        ADTypes.AutoZygote(),
+        [DI.gradient, DI.pullback],
+        scenarios;
+        correctness=ADTypes.AutoFiniteDifferences(
+            FiniteDifferences.central_fdm(5, 1)
+        ),
+        detailed=true,
+        rtol=1e-4,
+    )
 end
 
 @testset "Second derivatives" begin