Fix nesting of forwarddiff and sparsity detection

.github/workflows/ci.yml

@@ -21,7 +21,6 @@ jobs:
           - Downstream
         version:
           - '1'
-          - '1.6'
     steps:
       - uses: actions/checkout@v4
         with:

Project.toml

@@ -1,7 +1,7 @@
 name = "Symbolics"
 uuid = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 authors = ["Shashi Gowda <[email protected]>"]
-version = "5.14.1"
+version = "5.15.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -38,52 +38,56 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 
 [weakdeps]
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"
 
 [extensions]
+SymbolicsPreallocationToolsExt = "PreallocationTools"
 SymbolicsSymPyExt = "SymPy"
 
 [compat]
-ArrayInterface = "6, 7"
+ArrayInterface = "7"
 Bijections = "0.1"
-ConstructionBase = "1.1, 1.2"
+ConstructionBase = "1.2"
 DataStructures = "0.18"
 DiffRules = "1.4"
-Distributions = "0.23, 0.24, 0.25"
-DocStringExtensions = "0.7, 0.8, 0.9"
+Distributions = "0.25"
+DocStringExtensions = "0.9"
 DomainSets = "0.6"
 DynamicPolynomials = "0.5"
-Groebner = "0.3, 0.4, 0.5"
+Groebner = "0.5"
 IfElse = "0.1"
 LaTeXStrings = "1.3"
 LambertW = "0.4.5"
-Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
+Latexify = "0.16"
 LogExpFunctions = "0.3"
 MacroTools = "0.5"
-NaNMath = "0.3, 1"
+NaNMath = "1"
 PrecompileTools = "1"
 RecipesBase = "1.1"
-Reexport = "0.2, 1"
+Reexport = "1"
 ReferenceTests = "0.9"
 Requires = "1.1"
 RuntimeGeneratedFunctions = "0.5.9"
-SciMLBase = "1.8, 2"
-Setfield = "0.7, 0.8, 1"
-SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
+SciMLBase = "2"
+Setfield = "1"
+SpecialFunctions = "2"
 StaticArrays = "1.1"
 SymbolicIndexingInterface = "0.3"
 SymbolicUtils = "1.4"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "SafeTestsets", "Pkg", "PkgBenchmark", "BenchmarkTools", "ReferenceTests", "SymPy", "Random"]
+test = ["Test", "SafeTestsets", "Pkg", "PkgBenchmark", "PreallocationTools", "ForwardDiff", "BenchmarkTools", "ReferenceTests", "SymPy", "Random"]

ext/SymbolicsPreallocationToolsExt.jl

@@ -0,0 +1,49 @@
+module PreallocationToolsSymbolicsExt
+
+using PreallocationTools
+import PreallocationTools: _restructure, get_tmp
+using Symbolics, ForwardDiff
+
+function get_tmp(dc::DiffCache, u::Type{X}) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+function get_tmp(dc::DiffCache, u::X) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+function get_tmp(dc::DiffCache, u::AbstractArray{X}) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+function get_tmp(dc::FixedSizeDiffCache, u::Type{X}) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+function get_tmp(dc::FixedSizeDiffCache, u::X) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+function get_tmp(dc::FixedSizeDiffCache, u::AbstractArray{X}) where {T,N, X<: ForwardDiff.Dual{T, Num, N}}
+    if length(dc.du) > length(dc.any_du)
+        resize!(dc.any_du, length(dc.du))
+    end
+    _restructure(dc.du, dc.any_du)
+end
+
+end

test/nested_forwarddiff_sparsity.jl

@@ -0,0 +1,23 @@
+using ForwardDiff, SparseArrays, Symbolics, PreallocationTools
+# Test Nesting https://discourse.julialang.org/t/preallocationtools-jl-with-nested-forwarddiff-and-sparsity-pattern-detection-errors/107897
+
+function foo(x, cache)
+    d = get_tmp(cache, x)
+
+    d[:] = x
+
+    0.5 * x'*x
+end
+
+function residual(r, x, cache)
+    function foo_wrap(x)
+        foo(x, cache)
+    end
+
+    r[:] = ForwardDiff.gradient(foo_wrap, x)
+end
+
+cache = DiffCache(zeros(2))
+pattern = Symbolics.jacobian_sparsity((r, x) -> residual(r, x, cache), zeros(2), zeros(2))
+@test pattern == sparse([1 0
+                         0 1])

test/runtests.jl

@@ -34,6 +34,7 @@ if GROUP == "All" || GROUP == "Core"
     @safetestset "Algebraic Solver Test" begin include("solver.jl") end
     @safetestset "Groebner Bases Test" begin include("groebner_basis.jl") end
     @safetestset "Overloading Test" begin include("overloads.jl") end
+    @safetestset "Nested ForwardDiff Sparsity Test" begin include("nested_forwarddiff_sparsity.jl") end
     @safetestset "Build Function Test" begin include("build_function.jl") end
     @safetestset "Build Function Array Test" begin include("build_function_arrayofarray.jl") end
     @safetestset "Build Function Array Test Named Tuples" begin include("build_function_arrayofarray_named_tuples.jl") end