diff --git a/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml b/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml index ef04f99a6..f95a8351f 100644 --- a/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml +++ b/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml @@ -13,7 +13,7 @@ steps: agents: queue: "juliagpu" cuda: "*" - timeout_in_minutes: 30 + timeout_in_minutes: 120 # Don't run Buildkite if the commit message includes the text [skip tests] if: build.message !~ /\[skip tests\]/ diff --git a/lib/SimpleNonlinearSolve/Project.toml b/lib/SimpleNonlinearSolve/Project.toml index fdef9b8ee..4065b1fcf 100644 --- a/lib/SimpleNonlinearSolve/Project.toml +++ b/lib/SimpleNonlinearSolve/Project.toml @@ -1,7 +1,7 @@ name = "SimpleNonlinearSolve" uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" authors = ["SciML"] -version = "1.4.2" +version = "1.4.3" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/lib/SimpleNonlinearSolve/src/utils.jl b/lib/SimpleNonlinearSolve/src/utils.jl index 7390f5e3c..e89794b4c 100644 --- a/lib/SimpleNonlinearSolve/src/utils.jl +++ b/lib/SimpleNonlinearSolve/src/utils.jl @@ -40,17 +40,28 @@ function __pick_forwarddiff_chunk(x::StaticArray) end end -function __get_jacobian_config(ad::AutoForwardDiff{CS}, f, x) where {CS} +function __get_jacobian_config(ad::AutoForwardDiff{CS}, f::F, x) where {F, 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} +function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!::F, y, x) where {F, CS} ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}() tag = __standard_tag(ad.tag, x) return ForwardDiff.JacobianConfig(f!, y, x, ck, tag) end +function __forwarddiff_jacobian_config(f::F, x, ck::ForwardDiff.Chunk, tag) where {F} + return ForwardDiff.JacobianConfig(f, x, ck, tag) +end +function __forwarddiff_jacobian_config( + f::F, x::SArray, ck::ForwardDiff.Chunk{N}, tag) where {F, 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) :