Add AutoReactant{AutoEnzyme} (#127)

* Add AutoReactant{AutoEnzyme}

* Update index.md with new mode documentation

Added documentation for forward, reverse, or sparse mode.

* Update src/dense.jl

Co-authored-by: Avik Pal <avik.pal.2017@gmail.com>

* Fix documentation for AutoReactant mode parameter

* Update dense.jl

* Update parameterization description in AutoReactant

Clarified the description of the 'mode' field in AutoReactant.

* Apply suggestion from @gdalle

---------

Co-authored-by: Avik Pal <avik.pal.2017@gmail.com>
Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: 3 people

Project.toml

@@ -1,7 +1,7 @@
 name = "ADTypes"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com>, Guillaume Dalle and contributors"]
-version = "1.18.0"
+version = "1.19.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/index.md

@@ -55,6 +55,12 @@ AutoChainRules
 AutoDiffractor
 ```
 
+### Forward, reverse, or sparse mode
+
+```@docs
+AutoReactant{<:AutoEnzyme}
+```
+
 ### Symbolic mode
 
 ```@docs

src/ADTypes.jl

@@ -47,7 +47,8 @@ export AutoChainRules,
        AutoTracker,
        AutoZygote,
        NoAutoDiff,
-       NoAutoDiffSelectedError
+       NoAutoDiffSelectedError,
+       AutoReactant
 @public AbstractMode
 @public ForwardMode, ReverseMode, ForwardOrReverseMode, SymbolicMode
 @public mode

src/dense.jl

@@ -82,6 +82,45 @@ function Base.show(io::IO, backend::AutoEnzyme{M, A}) where {M, A}
     print(io, ")")
 end
 
+
+"""
+    AutoReactant{M<:AutoEnzyme}
+
+Struct used to select the [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) compilation atop Enzyme for automatic differentiation.
+
+Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
+
+# Constructors
+
+    AutoReactant(; mode::Union{AutoEnzyme,Nothing}=nothing)
+
+# Fields
+
+  - `mode::M` specifies the parameterization of differentiation.  It can be:
+
+      + an [`AutoEnzyme`](@ref) object if a specific mode is required
+      + `nothing` to choose the best mode automatically
+"""
+struct AutoReactant{M<:AutoEnzyme} <: AbstractADType
+    mode::M
+end
+
+function AutoReactant(;
+        mode::Union{AutoEnzyme,Nothing} = nothing)
+    if mode === nothing
+        mode = AutoEnzyme()
+    end
+    return AutoReactant(mode)
+end
+
+mode(r::AutoReactant) = mode(r.mode)
+
+function Base.show(io::IO, backend::AutoReactant)
+    print(io, AutoReactant, "(")
+    print(io, "mode=", repr(backend.mode; context = io))
+    print(io, ")")
+end
+
 """
     AutoFastDifferentiation
 

src/symbols.jl

@@ -22,7 +22,7 @@ ADTypes.AutoZygote()
 """
 Auto(package::Symbol, args...; kws...) = Auto(Val(package), args...; kws...)
 
-for backend in (:ChainRules, :Diffractor, :Enzyme, :FastDifferentiation,
+for backend in (:ChainRules, :Diffractor, :Enzyme, :Reactant, :FastDifferentiation,
     :FiniteDiff, :FiniteDifferences, :ForwardDiff, :Mooncake, :PolyesterForwardDiff,
     :ReverseDiff, :Symbolics, :Tapir, :Tracker, :Zygote)
     @eval Auto(::Val{$(QuoteNode(backend))}, args...; kws...) = $(Symbol(:Auto, backend))(

test/dense.jl

@@ -52,6 +52,34 @@ end
     @test ad.mode == EnzymeCore.Reverse
 end
 
+@testset "AutoReactant" begin
+    ad = AutoReactant()
+    @test ad isa AbstractADType
+    @test ad isa AutoReactant{<:AutoEnzyme}
+    @test ad.mode isa AutoEnzyme
+    @test ad.mode.mode === nothing
+    @test mode(ad) isa ForwardOrReverseMode
+
+    ad = AutoReactant(; mode=AutoEnzyme(; mode = EnzymeCore.Forward))
+    @test ad isa AbstractADType
+    @test ad isa AutoReactant{<:AutoEnzyme{typeof(EnzymeCore.Forward), Nothing}}
+    @test mode(ad) isa ForwardMode
+    @test ad.mode.mode == EnzymeCore.Forward
+
+    ad = AutoReactant(; mode=AutoEnzyme(; function_annotation = EnzymeCore.Const))
+    @test ad isa AbstractADType
+    @test ad isa AutoReactant{<:AutoEnzyme{Nothing, EnzymeCore.Const}}
+    @test mode(ad) isa ForwardOrReverseMode
+    @test ad.mode.mode === nothing
+
+    ad = AutoReactant(; mode=AutoEnzyme(;
+        mode = EnzymeCore.Reverse, function_annotation = EnzymeCore.Duplicated))
+    @test ad isa AbstractADType
+    @test ad isa AutoReactant{<:AutoEnzyme{typeof(EnzymeCore.Reverse), EnzymeCore.Duplicated}}
+    @test mode(ad) isa ReverseMode
+    @test ad.mode.mode == EnzymeCore.Reverse
+end
+
 @testset "AutoFastDifferentiation" begin
     ad = AutoFastDifferentiation()
     @test ad isa AbstractADType

test/symbols.jl

@@ -11,6 +11,7 @@ using Test
 @test ADTypes.Auto(:Mooncake) isa AutoMooncake
 @test ADTypes.Auto(:PolyesterForwardDiff) isa AutoPolyesterForwardDiff
 @test ADTypes.Auto(:ReverseDiff) isa AutoReverseDiff
+@test ADTypes.Auto(:Reactant) isa AutoReactant
 @test ADTypes.Auto(:Symbolics) isa AutoSymbolics
 @test ADTypes.Auto(:Tapir) isa AutoTapir
 @test ADTypes.Auto(:Tracker) isa AutoTracker