diff --git a/src/gradients.jl b/src/gradients.jl index 9458142..ae3a885 100644 --- a/src/gradients.jl +++ b/src/gradients.jl @@ -256,12 +256,46 @@ function finite_difference_gradient( dir = true) where {T1, T2, T3, T4, fdtype, returntype, inplace} if typeof(x) <: AbstractArray df = zero(returntype) .* x + finite_difference_gradient!( + df, f, x, cache, relstep = relstep, absstep = absstep, dir = dir) + df else - df = zero(cache.c1) + # Scalar x: compute out-of-place to support immutable output types + # (e.g. ArrayPartition{SVector} from SecondOrderODEProblem). + _scalar_gradient_oop(f, x, cache, fdtype, returntype, inplace; + relstep = relstep, absstep = absstep, dir = dir) + end +end + +# Out-of-place scalar→vector gradient that never mutates the result, +# so it works even when f returns immutable arrays (SVector, etc.). +function _scalar_gradient_oop( + f, x::Number, cache, fdtype, returntype, inplace; + relstep, absstep, dir) + fx, c1, c2 = cache.fx, cache.c1, cache.c2 + + if fdtype == Val(:forward) + epsilon = compute_epsilon(Val(:forward), x, relstep, absstep, dir) + _c1 = inplace == Val(true) ? (f(c1, x + epsilon); c1) : f(x + epsilon) + if typeof(fx) != Nothing + @. (_c1 - fx) / epsilon + else + _c2 = inplace == Val(true) ? (f(c2, x); c2) : f(x) + @. (_c1 - _c2) / epsilon + end + elseif fdtype == Val(:central) + epsilon = compute_epsilon(Val(:central), x, relstep, absstep, dir) + _c1 = inplace == Val(true) ? (f(c1, x + epsilon); c1) : f(x + epsilon) + _c2 = inplace == Val(true) ? (f(c2, x - epsilon); c2) : f(x - epsilon) + @. (_c1 - _c2) / (2 * epsilon) + elseif fdtype == Val(:complex) && returntype <: Real + epsilon_complex = eps(real(eltype(x))) + _c1 = inplace == Val(true) ? + (f(c1, x + im * epsilon_complex); c1) : f(x + im * epsilon_complex) + @. imag(_c1) / epsilon_complex + else + fdtype_error(returntype) end - finite_difference_gradient!( - df, f, x, cache, relstep = relstep, absstep = absstep, dir = dir) - df end # vector of derivatives of a vector->scalar map by each component of a vector x diff --git a/test/finitedifftests.jl b/test/finitedifftests.jl index 45d38e1..bba67be 100644 --- a/test/finitedifftests.jl +++ b/test/finitedifftests.jl @@ -290,6 +290,40 @@ complex_cache = FiniteDiff.GradientCache(df, x, Val{:complex}) @test err_func(FiniteDiff.finite_difference_gradient!(df, f, x, complex_cache), df_ref) < 1e-15 end +@time @testset "Gradient of f:scalar->vector with immutable output" begin + # Regression test: finite_difference_gradient with scalar x must work + # when f returns immutable arrays, since the out-of-place API should + # never need to mutate the result. This failed previously because the + # cached version called finite_difference_gradient! which used @. df = ... + # to write into an immutable buffer. + # + # We use a wrapper around a regular Vector that blocks setindex! to + # simulate immutable array types (like StaticArrays.SVector or + # ArrayPartition containing SVectors). + struct ReadOnlyVec{T} <: AbstractVector{T} + data::Vector{T} + end + Base.size(v::ReadOnlyVec) = size(v.data) + Base.getindex(v::ReadOnlyVec, i::Int) = v.data[i] + Base.setindex!(::ReadOnlyVec, _, ::Int) = error("ReadOnlyVec does not support setindex!") + Base.similar(v::ReadOnlyVec) = ReadOnlyVec(zeros(eltype(v), length(v))) + Base.zero(v::ReadOnlyVec) = ReadOnlyVec(zeros(eltype(v), length(v))) + # Out-of-place broadcast returns a plain Vector (like SVector .+ SVector returns SVector) + Base.BroadcastStyle(::Type{<:ReadOnlyVec}) = Broadcast.DefaultArrayStyle{1}() + + g(t) = ReadOnlyVec([sin(t), cos(t)]) + t0 = 1.0 + g_ref = [cos(t0), -sin(t0)] + + # Out-of-place cached version (the bug path) + df_template = similar(g(t0)) + for fdtype in (Val(:forward), Val(:central)) + cache = FiniteDiff.GradientCache(df_template, t0, fdtype, Float64, Val(false)) + result = FiniteDiff.finite_difference_gradient(g, t0, cache) + @test err_func(collect(result), g_ref) < 1e-4 + end +end + f(df, x) = (df[1] = sin(x); df[2] = cos(x); df) x = (2π * rand()) * (1 + im) fx = fill(zero(typeof(x)), 2)