Merge pull request #143 from JuliaDiff/jac_static

Specialize default chunk size on static arrays

Author: ChrisRackauckas

Project.toml

@@ -14,6 +14,7 @@ LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [compat]
@@ -25,8 +26,9 @@ FiniteDiff = "2"
 ForwardDiff = "0.10"
 LightGraphs = "1.3"
 Requires = "0.5, 1.0"
+StaticArrays = "1"
 VertexSafeGraphs = "0.1"
-julia = "1.4"
+julia = "1.6"
 
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"

src/SparseDiffTools.jl

@@ -12,6 +12,8 @@ using Adapt
 using LinearAlgebra
 using SparseArrays, ArrayInterface
 
+import StaticArrays
+
 using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using DataStructures: DisjointSets, find_root!, union!
 

src/differentiation/compute_jacobian_ad.jl

@@ -70,7 +70,15 @@ end
                 dx = sparsity === nothing && jac_prototype === nothing ? nothing : copy(x)) #if dx is nothing, we will estimate dx at the cost of a function call
 
     if sparsity === nothing && jac_prototype === nothing
-        cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
+        cfg = if chunksize === nothing
+            if typeof(x) <: StaticArrays.StaticArray
+                ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk{StaticArrays.Size(vec(x))[1]}())
+            else
+                ForwardDiff.JacobianConfig(f, x)
+            end
+        else
+            ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
+        end
         return ForwardDiff.jacobian(f, x, cfg)
     end
     if dx isa Nothing