Avoid MArrays with FiniteDiff backend

DifferentiationInterface/Project.toml

@@ -67,7 +67,7 @@ Diffractor = "=0.2.6"
 Enzyme = "0.13.39"
 EnzymeCore = "0.8.8"
 FastDifferentiation = "0.4.3"
-FiniteDiff = "2.27.0"
+FiniteDiff = "2.31.1"
 FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36,1"
 GPUArraysCore = "0.2"

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl

@@ -17,7 +17,7 @@ function DI.prepare_pushforward_nokwarg(
     cache = if x isa Number || y isa Number
         nothing
     else
-        JVPCache(similar(x), y, fdtype(backend))
+        JVPCache(copy(x), y, fdtype(backend))
     end
     relstep = if isnothing(backend.relstep)
         default_relstep(fdtype(backend), eltype(x))
@@ -133,7 +133,7 @@ function DI.prepare_derivative_nokwarg(
     cache = if y isa Number
         nothing
     elseif y isa AbstractArray
-        df = similar(y)
+        df = copy(y)
         cache = GradientCache(df, x, fdtype(backend), eltype(y), FUNCTION_NOT_INPLACE)
     end
     relstep = if isnothing(backend.relstep)
@@ -347,9 +347,9 @@ function DI.prepare_jacobian_nokwarg(
     _sig = DI.signature(f, backend, x, contexts...; strict)
     fc = DI.fix_tail(f, map(DI.unwrap, contexts)...)
     y = fc(x)
-    x1 = similar(x)
-    fx = similar(y)
-    fx1 = similar(y)
+    x1 = copy(x)
+    fx = copy(y)
+    fx1 = copy(y)
     cache = JacobianCache(x1, fx, fx1, fdjtype(backend))
     relstep = if isnothing(backend.relstep)
         default_relstep(fdjtype(backend), eltype(x))

DifferentiationInterface/test/Back/FiniteDiff/Project.toml

@@ -8,7 +8,8 @@ ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
-DifferentiationInterface = { path = "../../.." }
+DifferentiationInterface = {path = "../../.."}

DifferentiationInterface/test/Back/FiniteDiff/allocations.jl

@@ -0,0 +1,39 @@
+using StaticArrays: @SVector
+
+@testset "allocations checks" begin
+    function pushforward_allocs()
+        backend = AutoFiniteDiff()
+        x = @SVector [1.0, 2.0]
+        tx = (2.0 .* x,)
+        f(x) = @. 3.0 * x
+        prep = DifferentiationInterface.prepare_pushforward(f, backend, x, tx)
+        return prep
+    end
+    pushforward_allocs()
+    allocs = @allocated prep = pushforward_allocs()
+    # This needs FiniteDiff v2.31.1.
+    @test allocs == 0
+
+    function derivative_allocs()
+        backend = AutoFiniteDiff()
+        x = 3.0
+        f(x) = x .* (@SVector [1.0, 2.0])
+        prep = DifferentiationInterface.prepare_derivative(f, backend, x)
+        return prep
+    end
+    derivative_allocs()
+    allocs = @allocated prep = derivative_allocs()
+    @test allocs == 0
+
+    function jacobian_allocs()
+        backend = AutoFiniteDiff()
+        x = @SVector [1.0, 2.0]
+        f(x) = 3.0 .* x
+        prep = DifferentiationInterface.prepare_jacobian(f, backend, x)
+        return prep
+    end
+    jacobian_allocs()
+    allocs = @allocated prep = jacobian_allocs()
+    # Using FiniteDiff.jl with StaticArrays to calculate a Jacobian does result in some allocations, apparently because the `FiniteDiff.JacobianCache` is a `mutable struct`.
+    @test_broken allocs == 0
+end

DifferentiationInterface/test/Back/FiniteDiff/test.jl

@@ -113,3 +113,5 @@ end;
 end
 
 include("benchmark.jl")
+
+include("allocations.jl")