Put testing into subpackage (#110)

* Put testing into subpackage

* Fix deps

* Fix more deps

* More fixes

* Less reimporting

* Fix CI for docs

* Fix docs CI

* Rm JET

Author: gdalle

.github/workflows/CI.yml

@@ -62,6 +62,7 @@ jobs:
           using Pkg
           Pkg.Registry.update()  # TODO: remove? useful for fresh releases
           Pkg.develop(PackageSpec(path=pwd()))
+          Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib", "DifferentiationInterfaceTest")))
           Pkg.instantiate()
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-docdeploy@v1

Project.toml

@@ -7,58 +7,40 @@ version = "0.1.0"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
-JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Taped = "07d77754-e150-4737-8c94-cd238a1fb45b"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore"
-DifferentiationInterfaceChairmarksExt = ["Chairmarks"]
-DifferentiationInterfaceComponentArraysExt = ["ComponentArrays"]
 DifferentiationInterfaceDiffractorExt = ["AbstractDifferentiation", "Diffractor"]
 DifferentiationInterfaceEnzymeExt = "Enzyme"
-DifferentiationInterfaceFastDifferentiationExt = ["FastDifferentiation", "RuntimeGeneratedFunctions"]
+DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation"
 DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
-DifferentiationInterfaceForwardDiffExt = ["DiffResults", "ForwardDiff"]
-DifferentiationInterfaceJETExt = ["JET"]
-DifferentiationInterfaceJLArraysExt = ["JLArrays"]
-DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
-DifferentiationInterfaceReverseDiffExt = ["DiffResults", "ReverseDiff"]
-DifferentiationInterfaceStaticArraysExt = ["StaticArrays"]
-DifferentiationInterfaceTapedExt = ["Taped"]
-DifferentiationInterfaceTrackerExt = ["Tracker"]
-DifferentiationInterfaceZygoteExt = ["Zygote"]
+DifferentiationInterfaceForwardDiffExt = "ForwardDiff"
+DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff"
+DifferentiationInterfaceReverseDiffExt = "ReverseDiff"
+DifferentiationInterfaceTrackerExt = "Tracker"
+DifferentiationInterfaceZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.7"
 AbstractDifferentiation = "0.6"
 ChainRulesCore = "1.19"
-Chairmarks = "1.2"
-ComponentArrays = "0.15"
-DiffResults = "1.1"
 Diffractor = "0.2"
 DocStringExtensions = "0.9"
 Enzyme = "0.11"
@@ -67,16 +49,35 @@ FillArrays = "1"
 FiniteDiff = "2.22"
 FiniteDifferences = "0.12"
 ForwardDiff = "0.10"
-Functors = "0.4"
-JET = "0.8"
-JLArrays = "0.1"
 LinearAlgebra = "1"
 PolyesterForwardDiff = "0.1"
 ReverseDiff = "1.15"
-RuntimeGeneratedFunctions = "0.5"
-StaticArrays = "1.9"
-Taped = "1"
+Taped = "0.1"
 Test = "1"
 Tracker = "0.2"
 Zygote = "0.6"
 julia = "1.10"
+
+[extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[targets]
+test = ["ADTypes", "Aqua", "Chairmarks", "DataFrames", "Diffractor", "Documenter", "Enzyme", "FastDifferentiation", "FiniteDiff", "FiniteDifferences", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "PolyesterForwardDiff", "ReverseDiff", "Test", "Tracker", "Zygote"]

docs/Project.toml

@@ -4,7 +4,6 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
@@ -15,7 +14,6 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

docs/make.jl

@@ -1,10 +1,9 @@
 using Base: get_extension
 using DifferentiationInterface
-using DifferentiationInterface.DifferentiationTest
+using DifferentiationInterfaceTest
 import DifferentiationInterface as DI
 using Documenter
 using DocumenterMermaid
-using JET
 using Random
 using Test
 
@@ -40,7 +39,7 @@ makedocs(;
     modules=[
         ADTypes,
         DifferentiationInterface,
-        DifferentiationInterface.DifferentiationTest,
+        DifferentiationInterfaceTest,
         get_extension(DI, :DifferentiationInterfaceChainRulesCoreExt),
         get_extension(DI, :DifferentiationInterfaceDiffractorExt),
         get_extension(DI, :DifferentiationInterfaceEnzymeExt),

docs/src/api.md

@@ -1,5 +1,5 @@
 ```@meta
-CurrentModule = DifferentiationInterface
+CurrentModule = Main
 CollapsedDocStrings = true
 ```
 
@@ -61,7 +61,7 @@ Pages = ["prepare.jl"]
 ## Testing & benchmarking
 
 ```@autodocs
-Modules = [DifferentiationTest]
+Modules = [DifferentiationInterfaceTest]
 Private = false
 ```
 
@@ -77,6 +77,6 @@ Filter = t -> !(t isa Type && t <: ADTypes.AbstractADType)
 ```
 
 ```@autodocs
-Modules = [DifferentiationTest]
+Modules = [DifferentiationInterfaceTest]
 Public = false
 ```

docs/src/backends.md

@@ -1,10 +1,11 @@
 ```@meta
+CurrentModule = Main
 CollapsedDocStrings = true
 ```
 
 ```@setup backends
 using ADTypes, DifferentiationInterface
-using DifferentiationInterface.DifferentiationTest: backend_string
+using DifferentiationInterfaceTest: backend_string
 import Markdown
 import Enzyme, FastDifferentiation, FiniteDiff, FiniteDifferences, ForwardDiff, PolyesterForwardDiff, ReverseDiff, Tracker, Zygote
 

docs/src/tutorial.md

@@ -1,3 +1,7 @@
+```@meta
+CurrentModule = Main
+```
+
 # Tutorial
 
 We present a typical workflow with DifferentiationInterface.jl and showcase its potential performance benefits.
@@ -132,10 +136,10 @@ You didn't need to look at the docs of either ForwardDiff.jl or Enzyme.jl to ach
 ## Testing and benchmarking
 
 DifferentiationInterface.jl also provides some utilities for more involved comparison between backends.
-They are gathered in a submodule.
+They are gathered in a submodule called [`DifferentiationInterfaceTest`](https://github.com/gdalle/DifferentiationInterface.jl/tree/main/lib/DifferentiationInterfaceTest).
 
 ```@repl tuto
-using DifferentiationInterface.DifferentiationTest
+using DifferentiationInterfaceTest
 ```
 
 The main entry point is [`test_differentiation`](@ref), which is used as follows:

ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -1,7 +1,7 @@
 module DifferentiationInterfaceEnzymeExt
 
 using ADTypes: ADTypes, AutoEnzyme
-using DifferentiationInterface: mymul!!, myupdate!!, mysimilar, myzero, myzero!!
+using DifferentiationInterface: myupdate!!
 import DifferentiationInterface as DI
 using DocStringExtensions
 using Enzyme:
@@ -44,7 +44,7 @@ DI.mode(::AutoReverseEnzyme) = ADTypes.AbstractReverseMode
 
 # Enzyme's `Duplicated(x, dx)` expects both arguments to be of the same type
 function DI.basisarray(::AutoEnzyme, a::AbstractArray{T}, i::CartesianIndex) where {T}
-    b = myzero(a)
+    b = zero(a)
     b[i] = one(T)
     return b
 end

ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -1,25 +1,15 @@
 ## Pullback
 
-function DI.value_and_pullback!!(
-    f, _dx, ::AutoReverseEnzyme, x::Number, dy::Number, extras::Nothing
+function DI.value_and_pullback(
+    f, ::AutoReverseEnzyme, x::Number, dy::Number, extras::Nothing
 )
     der, y = autodiff(ReverseWithPrimal, f, Active, Active(x))
     new_dx = dy * only(der)
     return y, new_dx
 end
 
-function DI.value_and_pullback!!(
-    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, extras::Nothing
-)
-    dx_sametype = convert(typeof(x), dx)
-    dx_sametype = myzero!!(dx_sametype)
-    _, y = autodiff(ReverseWithPrimal, f, Active, Duplicated(x, dx_sametype))
-    dx_sametype = mymul!!(dx_sametype, dy)
-    return y, myupdate!!(dx, dx_sametype)
-end
-
-function DI.value_and_pullback!!(
-    f, _dx, ::AutoReverseEnzyme, x::Number, dy::AbstractArray, extras::Nothing
+function DI.value_and_pullback(
+    f, ::AutoReverseEnzyme, x::Number, dy::AbstractArray, extras::Nothing
 )
     forw, rev = autodiff_thunk(
         ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, Active{typeof(x)}
@@ -30,11 +20,21 @@ function DI.value_and_pullback!!(
     return y, new_dx
 end
 
+function DI.value_and_pullback!!(
+    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, extras::Nothing
+)
+    dx_sametype = convert(typeof(x), dx)
+    dx_sametype .= zero(eltype(dx_sametype))
+    _, y = autodiff(ReverseWithPrimal, f, Active, Duplicated(x, dx_sametype))
+    dx_sametype .*= dy
+    return y, myupdate!!(dx, dx_sametype)
+end
+
 function DI.value_and_pullback!!(
     f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::AbstractArray, extras::Nothing
 )
     dx_sametype = convert(typeof(x), dx)
-    dx_sametype = myzero!!(dx_sametype)
+    dx_sametype .= zero(eltype(dx_sametype))
     forw, rev = autodiff_thunk(
         ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, Duplicated{typeof(x)}
     )
@@ -44,25 +44,17 @@ function DI.value_and_pullback!!(
     return y, myupdate!!(dx, dx_sametype)
 end
 
-function DI.value_and_pullback(f, backend::AutoReverseEnzyme, x, dy, extras)
-    dx = mysimilar(x)
+function DI.value_and_pullback(f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras)
+    dx = similar(x)
     return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)
 end
 
 ## Gradient
 
-function DI.gradient(f, backend::AutoReverseEnzyme, x, extras::Nothing)
+function DI.gradient(f, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
     return gradient(Reverse, f, x)
 end
 
-function DI.gradient!!(f, grad, backend::AutoReverseEnzyme, x, extras::Nothing)
+function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
     return gradient!(Reverse, grad, f, x)
 end
-
-function DI.gradient(f, backend::AutoReverseEnzyme, x::Number, extras::Nothing)
-    return autodiff(Reverse, f, Active(x))
-end
-
-function DI.gradient!!(f, grad, backend::AutoReverseEnzyme, x::Number, extras::Nothing)
-    return autodiff(Reverse, f, Active(x))
-end

ext/DifferentiationInterfaceEnzymeExt/reverse_mutating.jl

@@ -8,10 +8,12 @@ function DI.value_and_pullback!!(
     return y, new_dx
 end
 
-function DI.value_and_pullback!!(f!, y, dx, ::AutoReverseEnzyme, x, dy, extras::Nothing)
+function DI.value_and_pullback!!(
+    f!, y, dx, ::AutoReverseEnzyme, x::AbstractArray, dy, extras::Nothing
+)
     dx_sametype = convert(typeof(x), dx)
-    dx_sametype = myzero!!(dx_sametype)
-    dy_sametype = convert(typeof(y), copy(dy))
+    dx_sametype .= zero(eltype(dx_sametype))
+    dy_sametype = convert(typeof(y), copy(dy))  # TODO: how to get rid of copy?
     autodiff(Reverse, f!, Const, Duplicated(y, dy_sametype), Duplicated(x, dx_sametype))
     return y, myupdate!!(dx, dx_sametype)
 end