Migrating to new ADTypes

Author: avik-pal

Project.toml

@@ -1,6 +1,6 @@
 name = "BoundaryValueDiffEq"
 uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
-version = "5.7.1"
+version = "5.8.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -24,6 +24,7 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 
@@ -34,7 +35,7 @@ ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
 BoundaryValueDiffEqODEInterfaceExt = "ODEInterface"
 
 [compat]
-ADTypes = "0.2.6"
+ADTypes = "1.2"
 Adapt = "4"
 Aqua = "0.8"
 ArrayInterface = "7.7"

src/BoundaryValueDiffEq.jl

@@ -192,7 +192,7 @@ end
                 nlsolve = NewtonRaphson(),
                 jac_alg = BVPJacobianAlgorithm(;
                     bc_diffmode = AutoForwardDiff(; chunksize = 2),
-                    nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))))
+                    nonbc_diffmode = AutoSparse(AutoForwardDiff(; chunksize = 2)))))
     end
 
     @compile_workload begin
@@ -257,13 +257,13 @@ end
                     nlsolve = LevenbergMarquardt(; disable_geodesic = Val(true)),
                     jac_alg = BVPJacobianAlgorithm(;
                         bc_diffmode = AutoForwardDiff(; chunksize = 2),
-                        nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
+                        nonbc_diffmode = AutoSparse(AutoForwardDiff(; chunksize = 2)))),
                 MultipleShooting(10,
                     Tsit5();
                     nlsolve = GaussNewton(),
                     jac_alg = BVPJacobianAlgorithm(;
                         bc_diffmode = AutoForwardDiff(; chunksize = 2),
-                        nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2)))])
+                        nonbc_diffmode = AutoSparse(AutoForwardDiff(; chunksize = 2))))])
     end
 
     @compile_workload begin

src/algorithms.jl

@@ -84,9 +84,9 @@ Significantly more stable than Single Shooting.
     on the input types and problem type.
 
       + For `TwoPointBVProblem`, only `diffmode` is used (defaults to
-        `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`).
+        `AutoSparse(AutoForwardDiff())` if possible else `AutoSparseFiniteDiff`).
       + For `BVProblem`, `bc_diffmode` and `nonbc_diffmode` are used. For `nonbc_diffmode`
-        we default to `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`. For
+        we default to `AutoSparse(AutoForwardDiff())` if possible else `AutoSparseFiniteDiff`. For
         `bc_diffmode`, we default to `AutoForwardDiff` if possible else `AutoFiniteDiff`.
   - `grid_coarsening`: Coarsening the multiple-shooting grid to generate a stable IVP
     solution. Possible Choices:
@@ -160,9 +160,9 @@ for order in (2, 3, 4, 5, 6)
             `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to
             use based on the input types and problem type.
             - For `TwoPointBVProblem`, only `diffmode` is used (defaults to
-              `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`).
+              `AutoSparse(AutoForwardDiff())` if possible else `AutoSparseFiniteDiff`).
             - For `BVProblem`, `bc_diffmode` and `nonbc_diffmode` are used. For
-              `nonbc_diffmode` defaults to `AutoSparseForwardDiff` if possible else
+              `nonbc_diffmode` defaults to `AutoSparse(AutoForwardDiff())` if possible else
               `AutoSparseFiniteDiff`. For `bc_diffmode`, defaults to `AutoForwardDiff` if
               possible else `AutoFiniteDiff`.
           - `defect_threshold`: Threshold for defect control.

src/solve/multiple_shooting.jl

@@ -113,7 +113,7 @@ function __solve_nlproblem!(
     jac_prototype = init_jacobian(jac_cache)
 
     ode_cache_jac_fn = __multiple_shooting_init_jacobian_odecache(
-        ensemblealg, prob, jac_cache, alg.jac_alg.diffmode,
+        ensemblealg, prob, jac_cache, __cache_trait(alg.jac_alg.diffmode),
         alg.ode_alg, cur_nshoot, u0; internal_ode_kwargs...)
 
     loss_fnₚ = @closure (du, u) -> __multiple_shooting_2point_loss!(
@@ -158,7 +158,7 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_
     ode_jac_cache = sparse_jacobian_cache(alg.jac_alg.nonbc_diffmode, sd_ode, nothing,
         similar(u_at_nodes, cur_nshoot * N), u_at_nodes)
     ode_cache_ode_jac_fn = __multiple_shooting_init_jacobian_odecache(
-        ensemblealg, prob, ode_jac_cache, alg.jac_alg.nonbc_diffmode,
+        ensemblealg, prob, ode_jac_cache, __cache_trait(alg.jac_alg.nonbc_diffmode),
         alg.ode_alg, cur_nshoot, u0; internal_ode_kwargs...)
 
     # BC Part
@@ -167,7 +167,7 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_
     bc_jac_cache = sparse_jacobian_cache(
         alg.jac_alg.bc_diffmode, sd_bc, nothing, similar(bcresid_prototype), u_at_nodes)
     ode_cache_bc_jac_fn = __multiple_shooting_init_jacobian_odecache(
-        ensemblealg, prob, bc_jac_cache, alg.jac_alg.bc_diffmode,
+        ensemblealg, prob, bc_jac_cache, __cache_trait(alg.jac_alg.bc_diffmode),
         alg.ode_alg, cur_nshoot, u0; internal_ode_kwargs...)
 
     jac_prototype = vcat(init_jacobian(bc_jac_cache), init_jacobian(ode_jac_cache))
@@ -208,12 +208,12 @@ function __multiple_shooting_init_odecache(
 end
 
 function __multiple_shooting_init_jacobian_odecache(
-        ensemblealg, prob, jac_cache, ad, alg, nshoots, u; kwargs...)
+        ensemblealg, prob, jac_cache, ::NoDiffCacheNeeded, alg, nshoots, u; kwargs...)
     return __multiple_shooting_init_odecache(ensemblealg, prob, alg, u, nshoots; kwargs...)
 end
 
-function __multiple_shooting_init_jacobian_odecache(ensemblealg, prob, jac_cache,
-        ::Union{AutoForwardDiff, AutoSparseForwardDiff}, alg, nshoots, u; kwargs...)
+function __multiple_shooting_init_jacobian_odecache(
+        ensemblealg, prob, jac_cache, ::DiffCacheNeeded, alg, nshoots, u; kwargs...)
     cache = jac_cache.cache
     if cache isa ForwardDiff.JacobianConfig
         xduals = reshape(cache.duals[2][1:length(u)], size(u))

src/solve/single_shooting.jl

@@ -49,7 +49,8 @@ function __solve(prob::BVProblem, alg_::Shooting; odesolve_kwargs = (;),
     end
 
     ode_cache_jac_fn = __single_shooting_jacobian_ode_cache(
-        internal_prob, jac_cache, alg.jac_alg.diffmode, u0, alg.ode_alg; ode_kwargs...)
+        internal_prob, jac_cache, __cache_trait(alg.jac_alg.diffmode),
+        u0, alg.ode_alg; ode_kwargs...)
 
     jac_prototype = init_jacobian(jac_cache)
 
@@ -126,13 +127,13 @@ function __single_shooting_jacobian(J, u, jac_cache, diffmode, loss_fn::L) where
     return J
 end
 
-function __single_shooting_jacobian_ode_cache(prob, jac_cache, alg, u0, ode_alg; kwargs...)
+function __single_shooting_jacobian_ode_cache(
+        prob, jac_cache, ::NoDiffCacheNeeded, u0, ode_alg; kwargs...)
     return SciMLBase.__init(remake(prob; u0), ode_alg; kwargs...)
 end
 
 function __single_shooting_jacobian_ode_cache(
-        prob, jac_cache, ::Union{AutoForwardDiff, AutoSparseForwardDiff},
-        u0, ode_alg; kwargs...)
+        prob, jac_cache, ::DiffCacheNeeded, u0, ode_alg; kwargs...)
     cache = jac_cache.cache
     if cache isa ForwardDiff.JacobianConfig
         xduals = cache.duals isa Tuple ? cache.duals[2] : cache.duals

src/types.jl

@@ -89,8 +89,8 @@ If user provided all the required fields, then return the user provided algorith
 Otherwise, based on the problem type and the algorithm, decide the missing fields.
 
 For example, for `TwoPointBVProblem`, the `bc_diffmode` is set to
-`AutoSparseForwardDiff` while for `StandardBVProblem`, the `bc_diffmode` is set to
-`AutoForwardDiff`.
+`AutoSparse(AutoForwardDiff())` while for `StandardBVProblem`, the `bc_diffmode` is set to
+`AutoForwardDiff()`.
 """
 function concrete_jacobian_algorithm(jac_alg::BVPJacobianAlgorithm, prob::BVProblem, alg)
     return concrete_jacobian_algorithm(jac_alg, prob.problem_type, prob, alg)
@@ -109,21 +109,13 @@ function concrete_jacobian_algorithm(
     return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode)
 end
 
-struct BoundaryValueDiffEqTag end
-
-function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:BoundaryValueDiffEqTag, <:T}},
-        f::F, x::AbstractArray{T}) where {T, F}
-    return true
-end
-
 @inline function __default_sparse_ad(x::AbstractArray{T}) where {T}
     return isbitstype(T) ? __default_sparse_ad(T) : __default_sparse_ad(first(x))
 end
 @inline __default_sparse_ad(x::T) where {T} = __default_sparse_ad(T)
 @inline __default_sparse_ad(::Type{<:Complex}) = AutoSparseFiniteDiff()
 @inline function __default_sparse_ad(::Type{T}) where {T}
-    return ForwardDiff.can_dual(T) ?
-           AutoSparseForwardDiff(; tag = BoundaryValueDiffEqTag()) : AutoSparseFiniteDiff()
+    return AutoSparse(ifelse(ForwardDiff.can_dual(T), AutoForwardDiff(), AutoFiniteDiff()))
 end
 
 @inline function __default_nonsparse_ad(x::AbstractArray{T}) where {T}
@@ -132,8 +124,7 @@ end
 @inline __default_nonsparse_ad(x::T) where {T} = __default_nonsparse_ad(T)
 @inline __default_nonsparse_ad(::Type{<:Complex}) = AutoFiniteDiff()
 @inline function __default_nonsparse_ad(::Type{T}) where {T}
-    return ForwardDiff.can_dual(T) ? AutoForwardDiff(; tag = BoundaryValueDiffEqTag()) :
-           AutoFiniteDiff()
+    return ifelse(ForwardDiff.can_dual(T), AutoForwardDiff(), AutoFiniteDiff())
 end
 
 # This can cause Type Instability
@@ -146,9 +137,10 @@ Base.@deprecate MIRKJacobianComputationAlgorithm(
     diffmode = missing; collocation_diffmode = missing, bc_diffmode = missing) BVPJacobianAlgorithm(
     diffmode; nonbc_diffmode = collocation_diffmode, bc_diffmode)
 
-__needs_diffcache(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = true
-__needs_diffcache(_) = false
-function __needs_diffcache(jac_alg::BVPJacobianAlgorithm)
+@inline __needs_diffcache(::AutoForwardDiff) = true
+@inline __needs_diffcache(ad::AutoSparse) = __needs_diffcache(ADTypes.dense_ad(ad))
+@inline __needs_diffcache(_) = false
+@inline function __needs_diffcache(jac_alg::BVPJacobianAlgorithm)
     return __needs_diffcache(jac_alg.diffmode) ||
            __needs_diffcache(jac_alg.bc_diffmode) ||
            __needs_diffcache(jac_alg.nonbc_diffmode)
@@ -176,3 +168,11 @@ const MaybeDiffCache = Union{DiffCache, FakeDiffCache}
         PreallocationTools.get_tmp(dc, u)
     end
 end
+
+# DiffCache
+struct DiffCacheNeeded end
+struct NoDiffCacheNeeded end
+
+@inline __cache_trait(::AutoForwardDiff) = DiffCacheNeeded()
+@inline __cache_trait(ad::AutoSparse) = __cache_trait(ADTypes.dense_ad(ad))
+@inline __cache_trait(_) = NoDiffCacheNeeded()

src/utils.jl

@@ -220,24 +220,8 @@ end
 __vec_bc(sol, p, t, bc, u_size) = vec(bc(__restructure_sol(sol, u_size), p, t))
 __vec_bc(sol, p, bc, u_size) = vec(bc(reshape(sol, u_size), p))
 
-__get_non_sparse_ad(ad::AbstractADType) = ad
-function __get_non_sparse_ad(ad::AbstractSparseADType)
-    if ad isa AutoSparseForwardDiff
-        return AutoForwardDiff{__get_chunksize(ad), typeof(ad.tag)}(ad.tag)
-    elseif ad isa AutoSparseEnzyme
-        return AutoEnzyme()
-    elseif ad isa AutoSparseFiniteDiff
-        return AutoFiniteDiff()
-    elseif ad isa AutoSparseReverseDiff
-        return AutoReverseDiff(ad.compile)
-    elseif ad isa AutoSparseZygote
-        return AutoZygote()
-    else
-        throw(ArgumentError("Unknown AD Type"))
-    end
-end
-
-__get_chunksize(::AutoSparseForwardDiff{CK}) where {CK} = CK
+@inline __get_non_sparse_ad(ad::AbstractADType) = ad
+@inline __get_non_sparse_ad(ad::AutoSparse) = ADTypes.dense_ad(ad)
 
 # Restructure Solution
 function __restructure_sol(sol::Vector{<:AbstractArray}, u_size)