From 110a076018f447675828e26bd2e27845e43367cb Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 29 Nov 2023 18:07:00 -0500
Subject: [PATCH 1/5] Non allocating versions for StaticArrays

---
 Project.toml                  |  2 +-
 src/SparseDiffTools.jl        |  2 +-
 src/highlevel/common.jl       | 62 ++++++++++++++++++++++++++++-------
 src/highlevel/forward_mode.jl |  8 +++++
 4 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6676a0ef..6583fd7c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
-version = "2.13.0"
+version = "2.14.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl
index 0d971507..238c9323 100644
--- a/src/SparseDiffTools.jl
+++ b/src/SparseDiffTools.jl
@@ -16,7 +16,7 @@ import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using ArrayInterface, SparseArrays
 import ArrayInterface: matrix_colors
 import StaticArrays
-import StaticArrays: StaticArray
+import StaticArrays: StaticArray, SArray, MArray, Size
 # Others
 using SciMLOperators, LinearAlgebra, Random
 import DataStructures: DisjointSets, find_root!, union!
diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl
index 8c8e3e51..688fee3c 100644
--- a/src/highlevel/common.jl
+++ b/src/highlevel/common.jl
@@ -173,6 +173,13 @@ A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`.
 """
 function sparse_jacobian_cache end
 
+function sparse_jacobian_static_array(ad, cache, f, x::SArray)
+    # Not the most performant fallback
+    J = init_jacobian(cache)
+    sparse_jacobian!(J, ad, cache, f, MArray(x))
+    return J
+end
+
 """
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f, x; fx=nothing)
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!, fx, x)
@@ -189,13 +196,18 @@ function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection,
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f,
+        x::SArray; kwargs...)
+    cache = sparse_jacobian_cache(ad, sd, f, x; kwargs...)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f, x)
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f!, fx, x)
 
 Use the sparsity detection `cache` for computing the sparse Jacobian. This allocates a new
-Jacobian at every function call
+Jacobian at every function call.
 """
 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache,
         args...)
@@ -203,6 +215,10 @@ function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianC
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f,
+        x::SArray)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType, sd::AbstractSparsityDetection,
@@ -247,14 +263,18 @@ function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) w
     C isa ForwardDiff.Chunk && return C
     return __chunksize(Val(C), x)
 end
-__chunksize(::Val{nothing}, x) = ForwardDiff.Chunk(x)
+__chunksize(::Val{nothing}, x) = __chunksize(x)
 function __chunksize(::Val{C}, x) where {C}
     if C isa Integer && !(C isa Bool)
-        return C ≤ 0 ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{C}()
+        return C ≤ 0 ? __chunksize(x) : ForwardDiff.Chunk{C}()
     else
         error("$(C)::$(typeof(C)) is not a valid chunksize!")
     end
 end
+
+__chunksize(x) = ForwardDiff.Chunk(x)
+__chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}()
+
 function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
     C === nothing && return nothing
     C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
@@ -273,18 +293,36 @@ end
     return :(nothing)
 end
 
-function init_jacobian(c::AbstractMaybeSparseJacobianCache)
+"""
+    init_jacobian(cache::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
+
+Initialize the Jacobian based on the cache. Uses sparse jacobians if possible.
+
+If `preserve_immutable` is `true`, then the Jacobian returned might be immutable, this is
+relevant if the inputs are immutable like `StaticArrays`.
+"""
+function init_jacobian(c::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
     T = promote_type(eltype(c.fx), eltype(c.x))
-    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x)
+    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x;
+        preserve_immutable)
 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
+function init_jacobian(::Nothing, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(fx, T, length(fx), length(x))
+end
+function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray;
+        preserve_immutable::Val{PI} = Val(true)) where {T, PI}
+    if PI
+        return StaticArrays.SArray{Tuple{length(fx), length(x)}, T}(I)
+    else
+        return StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
+    end
+end
+function init_jacobian(J, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(J, T, size(J, 1), size(J, 2))
 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)
+init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing
diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl
index 2ae378d6..08aec0d0 100644
--- a/src/highlevel/forward_mode.jl
+++ b/src/highlevel/forward_mode.jl
@@ -71,3 +71,11 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
     end
     return J
 end
+
+function sparse_jacobian_static_array(_, cache::ForwardDiffJacobianCache, f, x::SArray)
+    if cache.cache isa ForwardColorJacCache
+        return forwarddiff_color_jacobian(f, x, cache.cache)
+    else
+        return ForwardDiff.jacobian(f, x, cache.cache)
+    end
+end

From 69aaa0b93001c2996bc7532728fcc4a784ac1ea2 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 5 Dec 2023 10:29:16 -0500
Subject: [PATCH 2/5] Add tests

---
 Project.toml                 |  3 ++-
 src/highlevel/common.jl      | 17 +++++++++++++++++
 src/highlevel/finite_diff.jl |  2 ++
 test/test_sparse_jacobian.jl | 30 ++++++++++++++++++++++++++++--
 4 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6583fd7c..61066705 100644
--- a/Project.toml
+++ b/Project.toml
@@ -62,6 +62,7 @@ Zygote = "0.6"
 julia = "1.6"
 
 [extras]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -75,4 +76,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
+test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays", "AllocCheck"]
diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl
index 688fee3c..59098d98 100644
--- a/src/highlevel/common.jl
+++ b/src/highlevel/common.jl
@@ -188,6 +188,9 @@ Sequentially calls `sparse_jacobian_cache` and `sparse_jacobian!` to compute the
 `f` at `x`. Use this if the jacobian for `f` is computed exactly once. In all other
 cases, use `sparse_jacobian_cache` once to generate the cache and use `sparse_jacobian!`
 with the same cache to compute the jacobian.
+
+If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
+the jacobian computation. This is possible only for a limited backends currently.
 """
 function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, args...;
         kwargs...)
@@ -208,6 +211,9 @@ end
 
 Use the sparsity detection `cache` for computing the sparse Jacobian. This allocates a new
 Jacobian at every function call.
+
+If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
+the jacobian computation. This is possible only for a limited backends currently.
 """
 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache,
         args...)
@@ -326,3 +332,14 @@ init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing
+
+# Create a mutable version of the input array
+function __make_mutable(x)
+    if ArrayInterface.can_setindex(x)
+        return x
+    else
+        y = similar(x)
+        copyto!(y, x)
+        return y
+    end
+end
diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl
index 31114ae0..af8367ba 100644
--- a/src/highlevel/finite_diff.jl
+++ b/src/highlevel/finite_diff.jl
@@ -10,6 +10,7 @@ __getfield(c::FiniteDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype
 
 function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff},
         sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F}
+    x = __make_mutable(x) # FiniteDiff is bad at handling immutables
     coloring_result = sd(fd, f, x)
     fx = fx === nothing ? similar(f(x)) : fx
     if coloring_result isa NoMatrixColoring
@@ -25,6 +26,7 @@ end
 
 function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff},
         sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F}
+    x = __make_mutable(x) # FiniteDiff is bad at handling immutables
     coloring_result = sd(fd, f!, fx, x)
     if coloring_result isa NoMatrixColoring
         cache = FiniteDiff.JacobianCache(x, fx)
diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl
index 5b0a0f4a..3049df17 100644
--- a/test/test_sparse_jacobian.jl
+++ b/test/test_sparse_jacobian.jl
@@ -1,6 +1,6 @@
 ## Sparse Jacobian tests
-using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme
-using Test
+using AllocCheck, SparseDiffTools,
+    Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays
 
 @views function fdiff(y, x) # in-place
     L = length(x)
@@ -163,3 +163,29 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
         end
     end
 end
+
+# Testing that the non-sparse jacobian's are non-allocating.
+fvcat(x) = vcat(x, x)
+
+x_sa = @SVector randn(Float32, 10);
+
+J_true_sa = ForwardDiff.jacobian(fvcat, x_sa)
+
+@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
+    return sparse_jacobian(ad, sd, f, x)
+end
+
+@testset "Static Arrays" begin
+    @testset "No Allocations: $(difftype)" for difftype in (AutoSparseForwardDiff(),
+        AutoForwardDiff())
+        J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa)
+        @test J ≈ J_true_sa
+    end
+
+    @testset "Other Backends: $(difftype)" for difftype in (AutoSparseZygote(),
+        AutoZygote(), AutoSparseEnzyme(), AutoEnzyme(), AutoSparseFiniteDiff(),
+        AutoFiniteDiff())
+        J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa)
+        @test J ≈ J_true_sa
+    end
+end

From 2b67d665e043ffb0ffd243622331cc9ec3b399e5 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 5 Dec 2023 10:39:22 -0500
Subject: [PATCH 3/5] Conditionally add AllocCheck in 1.9

---
 Project.toml                 |  9 ++++----
 test/allocs/Project.toml     |  2 ++
 test/runtests.jl             |  7 +++---
 test/test_sparse_jacobian.jl | 42 ++++++++++++++++++++----------------
 4 files changed, 33 insertions(+), 27 deletions(-)
 create mode 100644 test/allocs/Project.toml

diff --git a/Project.toml b/Project.toml
index 61066705..d164d8a7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -45,13 +45,13 @@ Enzyme = "0.11"
 FiniteDiff = "2.8.1"
 ForwardDiff = "0.10"
 Graphs = "1"
-LinearAlgebra = "1.6"
+LinearAlgebra = "<0.0.1, 1"
 PackageExtensionCompat = "1"
-Random = "1.6"
+Random = "<0.0.1, 1"
 Reexport = "1"
 SciMLOperators = "0.3.7"
 Setfield = "1"
-SparseArrays = "1.6"
+SparseArrays = "<0.0.1, 1"
 StaticArrayInterface = "1.3"
 StaticArrays = "1"
 Symbolics = "5.5"
@@ -62,7 +62,6 @@ Zygote = "0.6"
 julia = "1.6"
 
 [extras]
-AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -76,4 +75,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays", "AllocCheck"]
+test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
diff --git a/test/allocs/Project.toml b/test/allocs/Project.toml
new file mode 100644
index 00000000..2ce39fac
--- /dev/null
+++ b/test/allocs/Project.toml
@@ -0,0 +1,2 @@
+[deps]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
diff --git a/test/runtests.jl b/test/runtests.jl
index 0708f4bc..52b53f92 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,8 +5,8 @@ const GROUP = get(ENV, "GROUP", "All")
 const is_APPVEYOR = (Sys.iswindows() && haskey(ENV, "APPVEYOR"))
 const is_TRAVIS = haskey(ENV, "TRAVIS")
 
-function activate_gpu_env()
-    Pkg.activate("gpu")
+function activate_env(env)
+    Pkg.activate(env)
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
     Pkg.instantiate()
 end
@@ -42,6 +42,7 @@ if GROUP == "Core" || GROUP == "All"
 end
 
 if GROUP == "InterfaceI" || GROUP == "All"
+    VERSION ≥ v"1.9" && activate_env("allocs")
     @time @safetestset "Jac Vecs and Hes Vecs" begin
         include("test_jaches_products.jl")
     end
@@ -54,7 +55,7 @@ if GROUP == "InterfaceI" || GROUP == "All"
 end
 
 if GROUP == "GPU"
-    activate_gpu_env()
+    activate_env("gpu")
     @time @safetestset "GPU AD" begin
         include("test_gpu_ad.jl")
     end
diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl
index 3049df17..1cccb01f 100644
--- a/test/test_sparse_jacobian.jl
+++ b/test/test_sparse_jacobian.jl
@@ -1,5 +1,5 @@
 ## Sparse Jacobian tests
-using AllocCheck, SparseDiffTools,
+using SparseDiffTools,
     Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays
 
 @views function fdiff(y, x) # in-place
@@ -164,28 +164,32 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
     end
 end
 
-# Testing that the non-sparse jacobian's are non-allocating.
-fvcat(x) = vcat(x, x)
+@static if VERSION ≥ v"1.9"
+    using AllocCheck
 
-x_sa = @SVector randn(Float32, 10);
+    # Testing that the non-sparse jacobian's are non-allocating.
+    fvcat(x) = vcat(x, x)
 
-J_true_sa = ForwardDiff.jacobian(fvcat, x_sa)
+    x_sa = @SVector randn(Float32, 10);
 
-@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
-    return sparse_jacobian(ad, sd, f, x)
-end
+    J_true_sa = ForwardDiff.jacobian(fvcat, x_sa)
 
-@testset "Static Arrays" begin
-    @testset "No Allocations: $(difftype)" for difftype in (AutoSparseForwardDiff(),
-        AutoForwardDiff())
-        J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa)
-        @test J ≈ J_true_sa
+    @check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
+        return sparse_jacobian(ad, sd, f, x)
     end
 
-    @testset "Other Backends: $(difftype)" for difftype in (AutoSparseZygote(),
-        AutoZygote(), AutoSparseEnzyme(), AutoEnzyme(), AutoSparseFiniteDiff(),
-        AutoFiniteDiff())
-        J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa)
-        @test J ≈ J_true_sa
+    @testset "Static Arrays" begin
+        @testset "No Allocations: $(difftype)" for difftype in (AutoSparseForwardDiff(),
+            AutoForwardDiff())
+            J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa)
+            @test J ≈ J_true_sa
+        end
+
+        @testset "Other Backends: $(difftype)" for difftype in (AutoSparseZygote(),
+            AutoZygote(), AutoSparseEnzyme(), AutoEnzyme(), AutoSparseFiniteDiff(),
+            AutoFiniteDiff())
+            J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa)
+            @test J ≈ J_true_sa
+        end
     end
-end
+end
\ No newline at end of file

From 55513d49fc202f76bc0be51059d71fc124786a9d Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 5 Dec 2023 10:51:26 -0500
Subject: [PATCH 4/5] Fixes the syntax?

---
 test/test_sparse_jacobian.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl
index 1cccb01f..d6d4273a 100644
--- a/test/test_sparse_jacobian.jl
+++ b/test/test_sparse_jacobian.jl
@@ -166,15 +166,17 @@ end
 
 @static if VERSION ≥ v"1.9"
     using AllocCheck
+end
 
+@static if VERSION ≥ v"1.9"
     # Testing that the non-sparse jacobian's are non-allocating.
     fvcat(x) = vcat(x, x)
 
-    x_sa = @SVector randn(Float32, 10);
+    x_sa = @SVector randn(Float32, 10)
 
     J_true_sa = ForwardDiff.jacobian(fvcat, x_sa)
 
-    @check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
+    AllocCheck.@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
         return sparse_jacobian(ad, sd, f, x)
     end
 
@@ -192,4 +194,4 @@ end
             @test J ≈ J_true_sa
         end
     end
-end
\ No newline at end of file
+end

From 7a8170fe3956f7f502c625b5896c1028fc383679 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 6 Dec 2023 10:48:20 -0500
Subject: [PATCH 5/5] Add direct SArray dispatch

---
 src/highlevel/common.jl      | 11 -----------
 src/highlevel/finite_diff.jl |  6 ++++--
 2 files changed, 4 insertions(+), 13 deletions(-)

diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl
index 59098d98..11ac4b8c 100644
--- a/src/highlevel/common.jl
+++ b/src/highlevel/common.jl
@@ -332,14 +332,3 @@ init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing
-
-# Create a mutable version of the input array
-function __make_mutable(x)
-    if ArrayInterface.can_setindex(x)
-        return x
-    else
-        y = similar(x)
-        copyto!(y, x)
-        return y
-    end
-end
diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl
index af8367ba..3a44258a 100644
--- a/src/highlevel/finite_diff.jl
+++ b/src/highlevel/finite_diff.jl
@@ -10,7 +10,6 @@ __getfield(c::FiniteDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype
 
 function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff},
         sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F}
-    x = __make_mutable(x) # FiniteDiff is bad at handling immutables
     coloring_result = sd(fd, f, x)
     fx = fx === nothing ? similar(f(x)) : fx
     if coloring_result isa NoMatrixColoring
@@ -26,7 +25,6 @@ end
 
 function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff},
         sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F}
-    x = __make_mutable(x) # FiniteDiff is bad at handling immutables
     coloring_result = sd(fd, f!, fx, x)
     if coloring_result isa NoMatrixColoring
         cache = FiniteDiff.JacobianCache(x, fx)
@@ -50,3 +48,7 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::FiniteDiffJacobianCache,
     FiniteDiff.finite_difference_jacobian!(J, f!, x, cache.cache)
     return J
 end
+
+function sparse_jacobian_static_array(_, cache::FiniteDiffJacobianCache, f, x::SArray)
+    return FiniteDiff.finite_difference_jacobian(f, x, cache.cache)
+end