Add tests

Author: avik-pal

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"]

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

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)

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