Remove support for autodiff::Symbol, autodiff::Bool, and chunk keyword arguments (#164)

* Remove support for `autodiff::Symbol`, `autodiff::Bool`, and `chunk` keyword arguments

* Use `DI.Cache`

Author: devmotion

Project.toml

@@ -1,12 +1,11 @@
 name = "NLSolversBase"
 uuid = "d41bc354-129a-5804-8e4c-c37616107c6c"
-version = "7.11.0"
+version = "8.0.0-dev"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
@@ -31,6 +30,7 @@ julia = "1.10"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 OptimTestProblems = "cec144fc-5a64-5bc6-99fb-dde8f63e154c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -40,4 +40,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "ComponentArrays", "ExplicitImports", "JET", "OptimTestProblems", "Random", "RecursiveArrayTools", "SparseArrays", "StaticArrays", "Test"]
+test = ["Aqua", "ComponentArrays", "ExplicitImports", "ForwardDiff", "JET", "OptimTestProblems", "Random", "RecursiveArrayTools", "SparseArrays", "StaticArrays", "Test"]

README.md

@@ -29,14 +29,7 @@ The words in front of `Differentiable` in the type names (`Non`, `Once`, `Twice`
 ## Automatic differentiation
 
 Some constructors for `OnceDifferentiable`, `TwiceDifferentiable`, `OnceDifferentiableConstraints` and `TwiceDifferentiableConstraints` accept a positional argument called `autodiff`.
-This argument can be either:
-
-- An object subtyping `AbstractADType`, defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl) and supported by [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl).
-- A `Symbol` like `:finite` (and variants thereof) or `:forward`, which fall back on `ADTypes.AutoFiniteDiff` and `ADTypes.AutoForwardDiff` respectively.
-- A `Bool`, namely `true`, which falls back on `ADTypes.AutoForwardDiff`.
-
-When the positional argument `chunk` is passed, it is used to configure chunk size in `ADTypes.AutoForwardDiff`, but _only_ if `autodiff in (:forward, true)`.
-Indeed, if `autodiff isa ADTypes.AutoForwardDiff`, we assume that the user already selected the appropriate chunk size and so `chunk` is ignored.
+This argument an instance of an `AbstractADType`, defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl) and supported by [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl).
 
 ## Examples
 #### Optimization

src/NLSolversBase.jl

@@ -1,9 +1,8 @@
 module NLSolversBase
 
-using ADTypes: AbstractADType, AutoForwardDiff, AutoFiniteDiff
+using ADTypes: AbstractADType, AutoFiniteDiff
 import DifferentiationInterface as DI
 using FiniteDiff: FiniteDiff
-using ForwardDiff: ForwardDiff
 using LinearAlgebra: LinearAlgebra
 
 export AbstractObjective,
@@ -45,36 +44,6 @@ export AbstractObjective,
 export AbstractConstraints, OnceDifferentiableConstraints,
     TwiceDifferentiableConstraints, ConstraintBounds
 
-function finitediff_fdtype(autodiff)
-    if autodiff == :finiteforward
-        return Val{:forward}
-    elseif autodiff == :finitecomplex
-        return Val{:complex}
-    elseif autodiff == :finite || autodiff == :central || autodiff == :finitecentral
-        return Val{:central}
-    else
-        throw(ArgumentError(LazyString("The autodiff value `", repr(autodiff), "` is not supported. Use `:finite` or `:forward`.")))
-    end
-end
-
-forwarddiff_chunksize(::Nothing) = nothing
-forwarddiff_chunksize(::ForwardDiff.Chunk{C}) where {C} = C
-
-is_finitediff(autodiff) = autodiff ∈ (:central, :finite, :finiteforward, :finitecomplex)
-is_forwarddiff(autodiff) = autodiff ∈ (:forward, :forwarddiff, true)
-
-get_adtype(autodiff::AbstractADType, chunk=nothing) = autodiff
-
-function get_adtype(autodiff::Union{Symbol,Bool}, chunk=nothing)
-    if is_finitediff(autodiff)
-        return AutoFiniteDiff(; fdtype=finitediff_fdtype(autodiff)())
-    elseif is_forwarddiff(autodiff)
-        return AutoForwardDiff(; chunksize=forwarddiff_chunksize(chunk))
-    else
-        throw(ArgumentError(LazyString("The autodiff value `", repr(autodiff), "` is not supported. Use `:finite` or `:forward`.")))
-    end
-end
-
 x_of_nans(x::AbstractArray, ::Type{Tf}=float(eltype(x))) where {Tf} = fill!(similar(x, Tf), NaN)
 
 include("objective_types/inplace_factory.jl")

src/objective_types/constraints.jl

@@ -119,18 +119,17 @@ function OnceDifferentiableConstraints(bounds::ConstraintBounds)
     OnceDifferentiableConstraints(c!, J!, bounds)
 end
 
-function checked_chunk(lx)
+function check_autodiff(lx::AbstractVector)
     if isempty(lx)
         throw(ArgumentError("autodiff on constraints require the full lower bound vector `lx`."))
     end
-    ForwardDiff.Chunk(lx)
 end
 
 function OnceDifferentiableConstraints(c!, lx::AbstractVector, ux::AbstractVector,
                                        lc::AbstractVector, uc::AbstractVector,
-                                       autodiff::Symbol = :central,
-                                       chunk::ForwardDiff.Chunk = checked_chunk(lx))
-                                       
+                                       autodiff::AbstractADType = AutoFiniteDiff(; fdtype = Val(:central)))
+
+    check_autodiff(lx)
     bounds = ConstraintBounds(lx, ux, lc, uc)
     T = eltype(bounds)
     sizex = size(lx)
@@ -139,11 +138,12 @@ function OnceDifferentiableConstraints(c!, lx::AbstractVector, ux::AbstractVecto
     xcache = zeros(T, sizex)
     ccache = zeros(T, sizec)
 
-    backend = get_adtype(autodiff, chunk)
-    jac_prep = DI.prepare_jacobian(c!, ccache, backend, xcache)
-    function j!(_j, _x)
-        DI.jacobian!(c!, ccache, _j, jac_prep, backend, _x)
-        return _j
+    jac_prep = DI.prepare_jacobian(c!, ccache, autodiff, xcache)
+    j! = let c! = c!, ccache = ccache, jac_prep = jac_prep, autodiff = autodiff
+        function (_j, _x)
+            DI.jacobian!(c!, ccache, _j, jac_prep, autodiff, _x)
+            return _j
+        end
     end
     return OnceDifferentiableConstraints(c!, j!, bounds)
 end
@@ -163,48 +163,45 @@ end
 
 function TwiceDifferentiableConstraints(c!, lx::AbstractVector, ux::AbstractVector,
     lc::AbstractVector, uc::AbstractVector,
-    autodiff::Symbol = :central,
-    chunk::ForwardDiff.Chunk = checked_chunk(lx))
+    autodiff::AbstractADType = AutoFiniteDiff(; fdtype = Val(:central)))
+    check_autodiff(lx)
     bounds = ConstraintBounds(lx, ux, lc, uc)
     T = eltype(bounds)
     nc = length(lc)
     nx = length(lx)
     x_example = zeros(T, nx)
     λ_example = zeros(T, nc)
     ccache = zeros(T, nc)
-   
-    function sum_constraints(_x, _λ)
-        # TODO: get rid of this allocation with DI.Cache
-        ccache_righttype = zeros(promote_type(T, eltype(_x)), nc)
-        c!(ccache_righttype, _x)
-        return LinearAlgebra.dot(_λ, ccache_righttype)
-    end
-
-    backend = get_adtype(autodiff, chunk)
-
 
-    jac_prep = DI.prepare_jacobian(c!, ccache, backend, x_example)
-    function con_jac!(_j, _x)
-        DI.jacobian!(c!, ccache, _j, jac_prep, backend, _x) 
-        return _j
+    jac_prep = DI.prepare_jacobian(c!, ccache, autodiff, x_example)
+    con_jac! = let c! = c!, ccache = ccache, jac_prep = jac_prep, autodiff = autodiff
+        function (_j, _x)
+            DI.jacobian!(c!, ccache, _j, jac_prep, autodiff, _x) 
+            return _j
+        end
     end
-
-    hess_prep = DI.prepare_hessian(sum_constraints, backend, x_example, DI.Constant(λ_example))
-    function con_hess!(_h, _x, _λ)
-        DI.hessian!(sum_constraints, _h, hess_prep, backend, _x, DI.Constant(_λ)) 
-        return _h
+    
+    function sum_constraints(_x, _λ, _ccache)
+        c!(_ccache, _x)
+        return LinearAlgebra.dot(_λ, _ccache)
+    end
+    hess_prep = DI.prepare_hessian(sum_constraints, autodiff, x_example, DI.Constant(λ_example), DI.Cache(ccache))
+    con_hess! = let sum_constraints = sum_constraints, hess_prep = hess_prep, autodiff = autodiff, ccache = ccache
+        function (_h, _x, _λ)
+            DI.hessian!(sum_constraints, _h, hess_prep, autodiff, _x, DI.Constant(_λ), DI.Cache(ccache)) 
+            return _h
+        end
     end
 
     return TwiceDifferentiableConstraints(c!, con_jac!, con_hess!, bounds) 
 end
 
 function TwiceDifferentiableConstraints(c!, con_jac!,lx::AbstractVector, ux::AbstractVector,
     lc::AbstractVector, uc::AbstractVector,
-    autodiff::Symbol = :central,
-    chunk::ForwardDiff.Chunk = checked_chunk(lx))
+    autodiff::AbstractADType = AutoFiniteDiff(; fdtype = Val(:central)),)
     # TODO: is con_jac! still useful? we ignore it here
 
-    return TwiceDifferentiableConstraints(c!, lx, ux, lc, uc, autodiff, chunk)
+    return TwiceDifferentiableConstraints(c!, lx, ux, lc, uc, autodiff)
 end
 
 

src/objective_types/oncedifferentiable.jl

@@ -13,26 +13,26 @@ end
 
 ### Only the objective
 # Ambiguity
-OnceDifferentiable(f, x::AbstractArray,
+function OnceDifferentiable(f, x::AbstractArray,
                    F::Real = real(zero(eltype(x))),
-                   DF::AbstractArray = alloc_DF(x, F); inplace::Bool = true, autodiff::Union{AbstractADType,Symbol,Bool} = :finite,  
-                   chunk::ForwardDiff.Chunk = ForwardDiff.Chunk(x)) =
-    OnceDifferentiable(f, x, F, DF, autodiff, chunk)
-#OnceDifferentiable(f, x::AbstractArray, F::AbstractArray; autodiff = :finite) =
-#    OnceDifferentiable(f, x::AbstractArray, F::AbstractArray, alloc_DF(x, F))
+                   DF::AbstractArray = alloc_DF(x, F);
+                   inplace::Bool = true,
+                   autodiff::AbstractADType = AutoFiniteDiff(; fdtype = Val(:central)))
+    OnceDifferentiable(f, x, F, DF, autodiff)
+end
 function OnceDifferentiable(f, x::AbstractArray,
                    F::AbstractArray, DF::AbstractArray = alloc_DF(x, F);
-                   inplace::Bool = true, autodiff::Union{AbstractADType,Symbol,Bool} = :finite)
+                   inplace::Bool = true,
+                   autodiff::AbstractADType = AutoFiniteDiff(; fdtype = Val(:central)))
     f! = f!_from_f(f, F, inplace)
-
     OnceDifferentiable(f!, x, F, DF, autodiff)
 end
 
 
 function OnceDifferentiable(f, x_seed::AbstractArray,
                             F::Real,
                             DF::AbstractArray,
-                            autodiff::Union{AbstractADType,Symbol,Bool}, chunk::ForwardDiff.Chunk)
+                            autodiff::AbstractADType)
     # When here, at the constructor with positional autodiff, it should already
     # be the case, that f is inplace.
     if f isa Union{InplaceObjective, NotInplaceObjective}
@@ -43,51 +43,48 @@ function OnceDifferentiable(f, x_seed::AbstractArray,
 
         return OnceDifferentiable(fF, dfF, fdfF, x_seed, F, DF)
     else
-        backend = get_adtype(autodiff, chunk)
-        grad_prep = DI.prepare_gradient(f, backend, x_seed)
-        function g!(_g, _x)
-            DI.gradient!(f, _g, grad_prep, backend, _x)
-            return nothing
+        grad_prep = DI.prepare_gradient(f, autodiff, x_seed)
+        g! = let f = f, grad_prep = grad_prep, autodiff = autodiff
+            function (_g, _x)
+                DI.gradient!(f, _g, grad_prep, autodiff, _x)
+                return nothing
+            end
         end
-        function fg!(_g, _x)
-            y, _ = DI.value_and_gradient!(f, _g, grad_prep, backend, _x)
-            return y
+        fg! = let f = f, grad_prep = grad_prep, autodiff = autodiff
+            function (_g, _x)
+                y, _ = DI.value_and_gradient!(f, _g, grad_prep, autodiff, _x)
+                return y
+            end
         end
         return OnceDifferentiable(f, g!, fg!, x_seed, F, DF)
     end
 end
 
-has_not_dep_symbol_in_ad = Ref{Bool}(true)
-OnceDifferentiable(f, x::AbstractArray, F::AbstractArray, autodiff::Symbol, chunk::ForwardDiff.Chunk = ForwardDiff.Chunk(x)) =
-OnceDifferentiable(f, x, F, alloc_DF(x, F), autodiff, chunk)
-function OnceDifferentiable(f, x::AbstractArray, F::AbstractArray,
-                            autodiff::Bool, chunk::ForwardDiff.Chunk = ForwardDiff.Chunk(x))
-    if autodiff == false
-        throw(ErrorException("It is not possible to set the `autodiff` keyword to `false` when constructing a OnceDifferentiable instance from only one function. Pass in the (partial) derivative or specify a valid `autodiff` symbol."))
-    elseif has_not_dep_symbol_in_ad[]
-        @warn("Setting the `autodiff` keyword to `true` is deprecated. Please use a valid symbol instead.")
-        has_not_dep_symbol_in_ad[] = false
-    end
-    OnceDifferentiable(f, x, F, alloc_DF(x, F), :forward, chunk)
-end
-function OnceDifferentiable(f, x_seed::AbstractArray, F::AbstractArray, DF::AbstractArray,
-    autodiff::Union{AbstractADType,Symbol,Bool}, chunk::ForwardDiff.Chunk = ForwardDiff.Chunk(x_seed))
-    if f isa Union{InplaceObjective, NotInplaceObjective}
+OnceDifferentiable(f, x::AbstractArray, F::AbstractArray, autodiff::AbstractADType) =
+    OnceDifferentiable(f, x, F, alloc_DF(x, F), autodiff)
+function OnceDifferentiable(f, x_seed::AbstractArray,
+                            F::AbstractArray,
+                            DF::AbstractArray,
+                            autodiff::AbstractADType)
+    if  f isa Union{InplaceObjective, NotInplaceObjective}
         fF = make_f(f, x_seed, F)
         dfF = make_df(f, x_seed, F)
         fdfF = make_fdf(f, x_seed, F)
         return OnceDifferentiable(fF, dfF, fdfF, x_seed, F, DF)
     else
         F2 = similar(F)
-        backend = get_adtype(autodiff, chunk)
-        jac_prep = DI.prepare_jacobian(f, F2, backend, x_seed)
-        function j!(_j, _x)
-            DI.jacobian!(f, F2, _j, jac_prep, backend, _x)
-            return _j
+        jac_prep = DI.prepare_jacobian(f, F2, autodiff, x_seed)
+        j! = let f = f, F2 = F2, jac_prep = jac_prep, autodiff = autodiff
+            function (_j, _x)
+                DI.jacobian!(f, F2, _j, jac_prep, autodiff, _x)
+                return _j
+            end
         end
-        function fj!(_y, _j, _x)
-            y, _ = DI.value_and_jacobian!(f, _y, _j, jac_prep, backend, _x)
-            return y
+        fj! = let f = f, jac_prep = jac_prep, autodiff = autodiff
+            function (_y, _j, _x)
+                y, _ = DI.value_and_jacobian!(f, _y, _j, jac_prep, autodiff, _x)
+                return y
+            end
         end
         return OnceDifferentiable(f, j!, fj!, x_seed, F, DF)
     end