diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl index b480578d0..698bc8465 100644 --- 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 index fc43543ec..fb4af1121 100644 --- 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 index 0c8737040..74667dbdb 100644 --- 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)