Allow a kwarg to force use the correct ordering

Author: avik-pal

src/differentiation/common.jl

@@ -38,32 +38,37 @@ __internal_oop(::JacFunctionWrapper{iip, oop}) where {iip, oop} = oop
 (f::JacFunctionWrapper{false, true, 2})(u) = f.f(u, f.p)
 (f::JacFunctionWrapper{false, true, 3})(u) = f.f(u)
 
-function JacFunctionWrapper(f::F, fu_, u, p, t) where {F}
+# NOTE: `use_deprecated_ordering` is a way for external libraries to update to the correct
+# style. In the next release, we will drop the first check
+function JacFunctionWrapper(f::F, fu_, u, p, t;
+        use_deprecated_ordering::Val{deporder} = Val(true)) where {F, deporder}
     # The warning instead of error ensures a non-breaking change for users relying on an
     # undefined / undocumented feature
     fu = fu_ === nothing ? copy(u) : copy(fu_)
 
-    # Check this first else we were breaking things
-    # In the next breaking release, we will fix the ordering of the checks
-    iip = static_hasmethod(f, typeof((fu, u)))
-    oop = static_hasmethod(f, typeof((u,)))
-    if iip || oop
-        if p !== nothing || t !== nothing
-            Base.depwarn("""`p` and/or `t` provided and are not `nothing`. But we
-            potentially detected `f(du, u)` or `f(u)`. This can be caused by:
+    if deporder
+        # Check this first else we were breaking things
+        # In the next breaking release, we will fix the ordering of the checks
+        iip = static_hasmethod(f, typeof((fu, u)))
+        oop = static_hasmethod(f, typeof((u,)))
+        if iip || oop
+            if p !== nothing || t !== nothing
+                Base.depwarn("""`p` and/or `t` provided and are not `nothing`. But we
+                potentially detected `f(du, u)` or `f(u)`. This can be caused by:
 
-            1. `f(du, u)` or `f(u)` is defined, in-which case `p` and/or `t` should not be
-               supplied.
-            2. `f(args...)` is defined, in which case `hasmethod` can be spurious.
+                1. `f(du, u)` or `f(u)` is defined, in-which case `p` and/or `t` should not
+                be supplied.
+                2. `f(args...)` is defined, in which case `hasmethod` can be spurious.
 
-            Currently, we perform the check for `f(du, u)` and `f(u)` first, but in future
-            breaking releases, this check will be performed last, which means that if `t`
-            is provided `f(du, u, p, t)`/`f(u, p, t)` will be given precedence, similarly
-            if `p` is provided `f(du, u, p)`/`f(u, p)` will be given precedence.""",
-                :JacFunctionWrapper)
+                Currently, we perform the check for `f(du, u)` and `f(u)` first, but in
+                future breaking releases, this check will be performed last, which means
+                that if `t` is provided `f(du, u, p, t)`/`f(u, p, t)` will be given
+                precedence, similarly if `p` is provided `f(du, u, p)`/`f(u, p)` will be
+                given precedence.""", :JacFunctionWrapper)
+            end
+            return JacFunctionWrapper{iip, oop, 3, F, typeof(fu), typeof(p), typeof(t)}(f,
+                fu, p, t)
         end
-        return JacFunctionWrapper{iip, oop, 3, F, typeof(fu), typeof(p), typeof(t)}(f,
-            fu, p, t)
     end
 
     if t !== nothing
@@ -86,6 +91,17 @@ function JacFunctionWrapper(f::F, fu_, u, p, t) where {F}
             fu, p, t)
     end
 
-    throw(ArgumentError("""Couldn't determine the function signature of `f` to construct a
-    JacobianWrapper!"""))
+    if !deporder
+        iip = static_hasmethod(f, typeof((fu, u)))
+        oop = static_hasmethod(f, typeof((u,)))
+        if !iip && !oop
+            throw(ArgumentError("""`p` is provided but `f(u)` or `f(fu, u)` not defined for
+            `f`!"""))
+        end
+        return JacFunctionWrapper{iip, oop, 3, F, typeof(fu), typeof(p), typeof(t)}(f,
+            fu, p, t)
+    else
+        throw(ArgumentError("""Couldn't determine the function signature of `f` to
+        construct a JacobianWrapper!"""))
+    end
 end

src/differentiation/jaches_products.jl

@@ -263,8 +263,9 @@ f(du, u)        # Otherwise
 ```
 """
 function JacVec(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing,
-        autodiff = AutoForwardDiff(), tag = DeivVecTag(), kwargs...)
-    ff = JacFunctionWrapper(f, fu, u, p, t)
+        autodiff = AutoForwardDiff(), tag = DeivVecTag(),
+        use_deprecated_ordering::Val = Val(true), kwargs...)
+    ff = JacFunctionWrapper(f, fu, u, p, t; use_deprecated_ordering)
     fu === nothing && (fu = __internal_oop(ff) ? ff(u) : u)
 
     cache, vecprod, vecprod! = if autodiff isa AutoFiniteDiff

src/differentiation/vecjac_products.jl

@@ -72,8 +72,8 @@ f(du, u)        # Otherwise
 ```
 """
 function VecJac(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing,
-        autodiff = AutoFiniteDiff(), kwargs...)
-    ff = JacFunctionWrapper(f, fu, u, p, t)
+        autodiff = AutoFiniteDiff(), use_deprecated_ordering::Val = Val(true), kwargs...)
+    ff = JacFunctionWrapper(f, fu, u, p, t; use_deprecated_ordering)
 
     if !__internal_oop(ff) && autodiff isa AutoZygote
         msg = "Zygote requires an out of place method with signature f(u)."