From aaff1efcd4554363c6242e95ada274488d492b7f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 14 Feb 2024 22:59:34 -0500
Subject: [PATCH] Static GPU compilation of Jacobian

---
 Project.toml |  2 +-
 src/utils.jl | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index bf94382..fdef9b8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.4.1"
+version = "1.4.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/utils.jl b/src/utils.jl
index 7390f5e..e353a02 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -43,7 +43,7 @@ end
 function __get_jacobian_config(ad::AutoForwardDiff{CS}, f, x) where {CS}
     ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
     tag = __standard_tag(ad.tag, x)
-    return ForwardDiff.JacobianConfig(f, x, ck, tag)
+    return __forwarddiff_jacobian_config(f, x, ck, tag)
 end
 function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!, y, x) where {CS}
     ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
@@ -51,6 +51,14 @@ function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!, y, x) where {CS}
     return ForwardDiff.JacobianConfig(f!, y, x, ck, tag)
 end
 
+__forwarddiff_jacobian_config(f, x, ck, tag) = ForwardDiff.JacobianConfig(f, x, ck, tag)
+function __forwarddiff_jacobian_config(f, x, ck::ForwardDiff.Chunk{N}, tag) where {N}
+    seeds = ForwardDiff.construct_seeds(ForwardDiff.Partials{N, eltype(x)})
+    duals = ForwardDiff.Dual{typeof(tag), eltype(x), N}.(x)
+    return ForwardDiff.JacobianConfig{typeof(tag), eltype(x), N, typeof(duals)}(seeds,
+        duals)
+end
+
 function __get_jacobian_config(ad::AutoPolyesterForwardDiff{CS}, args...) where {CS}
     x = last(args)
     return (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) :