From bb62c6c53db671b3de2d917983f170bce1f4cf7d Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 16 Jan 2024 09:47:35 -0500
Subject: [PATCH] Default to polyester forward diff

---
 Project.toml                                 | 10 +++++++---
 docs/src/tutorials/large_systems.md          |  1 +
 ext/NonlinearSolvePolyesterForwardDiffExt.jl |  7 +++++++
 src/internal/helpers.jl                      |  9 ++++++++-
 4 files changed, 23 insertions(+), 4 deletions(-)
 create mode 100644 ext/NonlinearSolvePolyesterForwardDiffExt.jl

diff --git a/Project.toml b/Project.toml
index 75f45bce7..15c99e73e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.5.0"
+version = "3.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -36,6 +36,7 @@ FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -48,6 +49,7 @@ NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
 NonlinearSolveNLsolveExt = "NLsolve"
+NonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSymbolicsExt = "Symbolics"
@@ -83,6 +85,7 @@ Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
+PolyesterForwardDiff = "0.1.1"
 Random = "1.91"
 RecursiveArrayTools = "3.4"
 Reexport = "1.2"
@@ -91,7 +94,7 @@ SafeTestsets = "0.1"
 SciMLBase = "2.19.0"
 SimpleNonlinearSolve = "1.2"
 SparseArrays = "1.10"
-SparseDiffTools = "2.14"
+SparseDiffTools = "2.16"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
@@ -121,6 +124,7 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -134,4 +138,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "PolyesterForwardDiff"]
diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md
index 38242c19f..4b9856b8b 100644
--- a/docs/src/tutorials/large_systems.md
+++ b/docs/src/tutorials/large_systems.md
@@ -61,6 +61,7 @@ The resulting `NonlinearProblem` definition is:
 
 ```@example ill_conditioned_nlprob
 using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, SparseDiffTools
+import PolyesterForwardDiff
 
 const N = 32
 const xyd_brusselator = range(0, stop = 1, length = N)
diff --git a/ext/NonlinearSolvePolyesterForwardDiffExt.jl b/ext/NonlinearSolvePolyesterForwardDiffExt.jl
new file mode 100644
index 000000000..52de8335f
--- /dev/null
+++ b/ext/NonlinearSolvePolyesterForwardDiffExt.jl
@@ -0,0 +1,7 @@
+module NonlinearSolvePolyesterForwardDiffExt
+
+using NonlinearSolve, PolyesterForwardDiff
+
+NonlinearSolve.is_extension_loaded(::Val{:PolyesterForwardDiff}) = true
+
+end
diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl
index f9e90b7f8..b70331f18 100644
--- a/src/internal/helpers.jl
+++ b/src/internal/helpers.jl
@@ -63,7 +63,14 @@ function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, ar
         use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
     else
         tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))
-        (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; tag)
+        if use_sparse_ad
+            AutoSparseForwardDiff(; tag) # Polyester Sparse Mode is not implemented yet
+        elseif is_extension_loaded(Val(:PolyesterForwardDiff)) &&
+               !(prob.u0 isa Number || prob.u0 isa SArray)
+            AutoPolyesterForwardDiff()
+        else
+            AutoForwardDiff(; tag)
+        end
     end
     return ad
 end