Introduce AutoDI wrapper type for DifferentiationInterface dispatch

src/ADTypes.jl

@@ -58,6 +58,10 @@ export AutoChainRules,
 export AutoSparse
 @public dense_ad
 
+# DI Automatic Differentiation
+export AutoDI
+@public inner_ad
+
 # Sparsity detection
 export AbstractSparsityDetector
 export jacobian_sparsity, hessian_sparsity

src/dense.jl

@@ -586,3 +586,53 @@ NoAutoDiffSelectedError() = NoAutoDiffSelectedError("Automatic differentiation c
 function mode(::NoAutoDiff)
     throw(NoAutoDiffSelectedError())
 end
+
+"""
+    AutoDI{I<:AbstractADType}
+
+Wraps an AD type to signify that the DifferentiationInterface wrapper should be used instead of calling the backend directly.
+
+This allows packages to distinguish between an intention to directly call a corresponding AD tool vs. the DI wrapper for said tool, enabling the ability to use, test, and validate both approaches.
+
+# Fields
+
+  - `inner_ad::I`: the underlying AD package, subtyping [`AbstractADType`](@ref)
+
+# Constructors
+
+    AutoDI(inner_ad)
+
+# Example
+
+```jldoctest
+julia> using ADTypes
+
+julia> ad = AutoDI(AutoForwardDiff())
+AutoDI(AutoForwardDiff())
+
+julia> inner_ad(ad)
+AutoForwardDiff()
+```
+"""
+struct AutoDI{I <: AbstractADType} <: AbstractADType
+    inner_ad::I
+end
+
+function Base.show(io::IO, backend::AutoDI)
+    print(io, AutoDI, "(", repr(backend.inner_ad, context = io), ")")
+end
+
+"""
+    inner_ad(ad::AutoDI)::AbstractADType
+    inner_ad(ad::AbstractADType)::AbstractADType
+
+Return the underlying AD package for a DI AD choice, acts as the identity on a non-DI AD choice.
+
+# See also
+
+  - [`AutoDI`](@ref)
+"""
+inner_ad(ad::AutoDI) = ad.inner_ad
+inner_ad(ad::AbstractADType) = ad
+
+mode(di_ad::AutoDI) = mode(inner_ad(di_ad))

test/di.jl

@@ -0,0 +1,65 @@
+@testset "AutoDI" begin
+    @testset "Subtyping and wrapping $ad_name" for (ad_name, ad) in [
+        ("AutoForwardDiff", AutoForwardDiff()),
+        ("AutoZygote", AutoZygote()),
+        ("AutoEnzyme", AutoEnzyme()),
+        ("AutoReverseDiff", AutoReverseDiff()),
+        ("AutoChainRules", AutoChainRules(; ruleconfig = ForwardOrReverseRuleConfig())),
+    ]
+        di_ad = AutoDI(ad)
+        @test di_ad isa AbstractADType
+        @test di_ad isa AutoDI
+
+        # Test mode propagation
+        if mode(ad) isa ForwardMode
+            @test mode(di_ad) isa ForwardMode
+        elseif mode(ad) isa ForwardOrReverseMode
+            @test mode(di_ad) isa ForwardOrReverseMode
+        elseif mode(ad) isa ReverseMode
+            @test mode(di_ad) isa ReverseMode
+        elseif mode(ad) isa SymbolicMode
+            @test mode(di_ad) isa SymbolicMode
+        end
+
+        # Test inner_ad accessor
+        @test inner_ad(ad) == ad
+        @test inner_ad(di_ad) == ad
+    end
+
+    @testset "All AD backends" begin
+        for ad in every_ad()
+            di_ad = AutoDI(ad)
+            @test di_ad isa AbstractADType
+            @test inner_ad(di_ad) == ad
+            @test mode(di_ad) == mode(ad)
+        end
+    end
+
+    @testset "Nested wrapping" begin
+        # Test that we can wrap AutoDI with AutoSparse and vice versa
+        ad = AutoForwardDiff()
+        di_ad = AutoDI(ad)
+        sparse_di_ad = AutoSparse(di_ad)
+
+        @test sparse_di_ad isa AutoSparse
+        @test dense_ad(sparse_di_ad) isa AutoDI
+        @test inner_ad(dense_ad(sparse_di_ad)) == ad
+
+        # Test AutoDI wrapping AutoSparse
+        sparse_ad = AutoSparse(ad)
+        di_sparse_ad = AutoDI(sparse_ad)
+
+        @test di_sparse_ad isa AutoDI
+        @test inner_ad(di_sparse_ad) isa AutoSparse
+        @test dense_ad(inner_ad(di_sparse_ad)) == ad
+    end
+
+    @testset "Display" begin
+        ad = AutoForwardDiff(chunksize = 5)
+        di_ad = AutoDI(ad)
+
+        str = sprint(show, di_ad)
+        @test occursin("AutoDI", str)
+        @test occursin("AutoForwardDiff", str)
+    end
+end

test/runtests.jl

@@ -6,6 +6,7 @@ using ADTypes: AbstractADType,
                ReverseMode,
                SymbolicMode
 using ADTypes: dense_ad,
+               inner_ad,
                NoSparsityDetector,
                KnownJacobianSparsityDetector,
                KnownHessianSparsityDetector,
@@ -101,6 +102,9 @@ end
     @testset "Sparse" begin
         include("sparse.jl")
     end
+    @testset "DI" begin
+        include("di.jl")
+    end
     @testset "Symbols" begin
         include("symbols.jl")
     end