Don't use custom tags

SciML · Apr 16, 2024 · 316f74e · 316f74e
1 parent e295922
commit 316f74e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 12 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.10.0"
+version = "3.10.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

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
@@ -103,9 +103,9 @@ 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))),
+                cache1 = Dual{typeof(ForwardDiff.Tag(uf, eltype(u))),
                     eltype(u), 1}.(similar(u), ForwardDiff.Partials.(tuple.(u)))
-                cache2 = Dual{typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(fu))),
+                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