Special Dispatch for Static Array

JuliaDiff · Oct 8, 2023 · cf98e97 · cf98e97
1 parent 659b523
commit cf98e97
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl
@@ -16,6 +16,7 @@ import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using ArrayInterface, SparseArrays
 import ArrayInterface: matrix_colors
 import StaticArrays
+import StaticArrays: StaticArray
 # Others
 using SciMLOperators, LinearAlgebra, Random
 import DataStructures: DisjointSets, find_root!, union!

diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl
@@ -259,6 +259,11 @@ function init_jacobian(c::AbstractMaybeSparseJacobianCache)
     return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x)
 end
 init_jacobian(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x))
+function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray) where {T}
+    # We want to construct a MArray to preserve types
+    J = StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
+    return J
+end
 init_jacobian(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2))
 init_jacobian(J::SparseMatrixCSC, ::Type{T}, _, _) where {T} = T.(J)