Merge pull request #411 from SciML/ap/untag

Don't use custom tags
SciML · Apr 18, 2024 · b389d0e · b389d0e
2 parents 9a0bb8d + c60a2b6
commit b389d0e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 14 deletions.
diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl
@@ -34,7 +34,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::NLSolversJL, args...;
 
         if autodiff_concrete === :forwarddiff
             fj_scalar = @closure (Jx, x) -> begin
-                T = typeof(NonlinearSolve.NonlinearSolveTag())
+                T = typeof(ForwardDiff.Tag(prob.f, eltype(x)))
                 x_dual = ForwardDiff.Dual{T}(x, one(x))
                 y = prob.f(x_dual, prob.p)
                 return ForwardDiff.value(y), ForwardDiff.extract_derivative(T, y)

diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl
@@ -30,13 +30,6 @@ function evaluate_f!!(f::NonlinearFunction{iip}, fu, u, p) where {iip}
 end
 
 # AutoDiff Selection Functions
-struct NonlinearSolveTag end
-
-function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T}},
-        f::F, x::AbstractArray{T}) where {T, F}
-    return true
-end
-
 function get_concrete_forward_ad(
         autodiff::Union{ADTypes.AbstractForwardMode, ADTypes.AbstractFiniteDifferencesMode},
         prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse}
@@ -62,8 +55,7 @@ function get_concrete_forward_ad(
     ad = if !ForwardDiff.can_dual(eltype(prob.u0)) # Use Finite Differencing
         use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
     else
-        tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))
-        (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; tag)
+        (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)()
     end
     return ad
 end

diff --git a/src/internal/operators.jl b/src/internal/operators.jl
@@ -109,10 +109,10 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
         if jvp_autodiff isa AutoForwardDiff || jvp_autodiff isa AutoPolyesterForwardDiff
             if iip
                 # FIXME: Technically we should propagate the tag but ignoring that for now
-                cache1 = Dual{typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(u))),
-                    eltype(u), 1}.(similar(u), ForwardDiff.Partials.(tuple.(u)))
-                cache2 = Dual{typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(fu))),
-                    eltype(fu), 1}.(similar(fu), ForwardDiff.Partials.(tuple.(fu)))
+                cache1 = Dual{typeof(ForwardDiff.Tag(uf, eltype(u))), eltype(u),
+                    1}.(similar(u), ForwardDiff.Partials.(tuple.(u)))
+                cache2 = Dual{typeof(ForwardDiff.Tag(uf, eltype(fu))), eltype(fu),
+                    1}.(similar(fu), ForwardDiff.Partials.(tuple.(fu)))
                 @closure (Jv, v, u, p) -> auto_jacvec!(Jv, uf, u, v, cache1, cache2)
             else
                 @closure (v, u, p) -> auto_jacvec(uf, u, v)