From 8ea0814a85442ffdb4aa744b5addb3d3fcb3c4f5 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 3 Aug 2026 12:33:56 -0400 Subject: [PATCH] Make the out-of-place jacobian's return type inferrable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `finite_difference_jacobian` (out-of-place) inferred to `Any`. Its per-color closures assign variables that also exist in the enclosing function's scope — `epsilon`, `vecfx`, `x_save`, `dx`, ... — so those became captured variables with several assignment sites, which closure conversion boxes. Boxed captures read back as `Any`, and that `Any` propagates through `mapreduce(calculate_Ji_*, hcat, ...)` all the way to the returned matrix. Declare the closures' temporaries `local`, rename the ones that had to stay distinct from a capture, and give `vecfx`/`rows_index`/`cols_index` a single assignment site each. No behavioral change: the closures only ever wrote those outer variables by accident, and no branch read the values back. This surfaced downstream as `@inferred SciMLBase.init(prob, alg)` failing for a `TrustRegion(autodiff = AutoFiniteDiff())` `NonlinearProblem`, whose Jacobian cache is typed by this return value (SciML/NonlinearSolve.jl#1092). Static-array inputs remain uninferrable: the `hcat` accumulator gains an `SMatrix` column per color, so its type depends on `maximum(colorvec)`. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01MoeiVWqAEFJQK22kiXu48f --- Project.toml | 2 +- src/jacobians.jl | 52 ++++++++++++++++++++++---------------- test/out_of_place_tests.jl | 14 ++++++++++ 3 files changed, 45 insertions(+), 23 deletions(-) diff --git a/Project.toml b/Project.toml index 479b304..f415eed 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FiniteDiff" uuid = "6a86dc24-6348-571c-b903-95158fe2bd41" -version = "2.32.0" +version = "2.32.1" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" diff --git a/src/jacobians.jl b/src/jacobians.jl index 108a60d..7a003e3 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -294,14 +294,17 @@ function finite_difference_jacobian( copyto!(x1, x) end - if !(f_in isa Nothing) - vecfx = _vec(f_in) + # Single assignment site: `vecfx` is captured by the `calculate_Ji_*` closures + # below, and a captured variable with several assignment sites is boxed, which + # makes the whole function's return type uninferrable. + vecfx = if !(f_in isa Nothing) + _vec(f_in) elseif fdtype == Val(:forward) - vecfx = _vec(f(x)) + _vec(f(x)) elseif fdtype == Val(:complex) && returntype <: Real - vecfx = real(fx) + real(fx) else - vecfx = _vec(fx) + _vec(fx) end vecx = _vec(x) J = jac_prototype isa Nothing ? @@ -310,13 +313,16 @@ function finite_difference_jacobian( nrows, ncols = size(J) if !(sparsity isa Nothing) - rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) - rows_index = [rows_index[i] for i in 1:length(rows_index)] - cols_index = [cols_index[i] for i in 1:length(cols_index)] + structural_rows, structural_cols = ArrayInterface.findstructralnz(sparsity) + rows_index = [structural_rows[i] for i in 1:length(structural_rows)] + cols_index = [structural_cols[i] for i in 1:length(structural_cols)] end if fdtype == Val(:forward) + # `local` keeps these from aliasing (and thereby boxing) the same-named + # variables of the enclosing function; see the `vecfx` comment above. function calculate_Ji_forward(i) + local x_save, epsilon, _vecx1, _x1, vecfx1, dx x_save = ArrayInterface.allowed_getindex(vecx, i) epsilon = compute_epsilon(Val(:forward), x_save, relstep, absstep, dir) _vecx1 = setindex(vecx, x_save+epsilon, i) @@ -336,12 +342,12 @@ function finite_difference_jacobian( J = J + _make_Ji(J, eltype(x), dx, color_i, nrows, ncols) else tmp = norm(vecx .* (colorvec .== color_i)) - epsilon = compute_epsilon( + epsilon_c = compute_epsilon( Val(:forward), sqrt(tmp), relstep, absstep, dir) - _vecx = @. vecx + epsilon * (colorvec == color_i) + _vecx = @. vecx + epsilon_c * (colorvec == color_i) _x = reshape(_vecx, axes(x)) vecfx1 = _vec(f(_x)) - dx = (vecfx1-vecfx)/epsilon + dx = (vecfx1-vecfx)/epsilon_c Ji = _make_Ji( J, rows_index, cols_index, dx, colorvec, color_i, nrows, ncols) J = J + Ji @@ -354,6 +360,7 @@ function finite_difference_jacobian( # unsafe — the cache may have been built via `similar(x)` or reused at a # different x — so we always perturb around `vecx` directly. function calculate_Ji_central(i) + local x_save, epsilon, _vecx1, _vecx, _x1, _x, vecfx1, vecfx0, dx x_save = ArrayInterface.allowed_getindex(vecx, i) epsilon = compute_epsilon(Val(:forward), x_save, relstep, absstep, dir) _vecx1 = setindex(vecx, x_save+epsilon, i) @@ -361,8 +368,8 @@ function finite_difference_jacobian( _x1 = reshape(_vecx1, axes(x)) _x = reshape(_vecx, axes(x)) vecfx1 = _vec(f(_x1)) - vecfx = _vec(f(_x)) - dx = (vecfx1-vecfx)/(2epsilon) + vecfx0 = _vec(f(_x)) + dx = (vecfx1-vecfx0)/(2epsilon) return dx end @@ -376,15 +383,15 @@ function finite_difference_jacobian( J = J + _make_Ji(J, eltype(x), dx, color_i, nrows, ncols) else tmp = norm(vecx .* (colorvec .== color_i)) - epsilon = compute_epsilon( + epsilon_c = compute_epsilon( Val(:forward), sqrt(tmp), relstep, absstep, dir) - _vecx1 = @. vecx + epsilon * (colorvec == color_i) - _vecx = @. vecx - epsilon * (colorvec == color_i) + _vecx1 = @. vecx + epsilon_c * (colorvec == color_i) + _vecx = @. vecx - epsilon_c * (colorvec == color_i) _x1 = reshape(_vecx1, axes(x)) _x = reshape(_vecx, axes(x)) vecfx1 = _vec(f(_x1)) - vecfx = _vec(f(_x)) - dx = (vecfx1-vecfx)/(2epsilon) + vecfx0 = _vec(f(_x)) + dx = (vecfx1-vecfx0)/(2epsilon_c) Ji = _make_Ji( J, rows_index, cols_index, dx, colorvec, color_i, nrows, ncols) J = J + Ji @@ -395,11 +402,12 @@ function finite_difference_jacobian( epsilon = eps(eltype(x)) function calculate_Ji_complex(i) + local x_save, _vecx, _x, vecfx_c, dx x_save = ArrayInterface.allowed_getindex(vecx, i) _vecx = setindex(complex.(vecx), x_save+im*epsilon, i) _x = reshape(_vecx, axes(x)) - vecfx = _vec(f(_x)) - dx = imag(vecfx)/epsilon + vecfx_c = _vec(f(_x)) + dx = imag(vecfx_c)/epsilon return dx end @@ -414,8 +422,8 @@ function finite_difference_jacobian( else _vecx = @. vecx + im * epsilon * (colorvec == color_i) _x = reshape(_vecx, axes(x)) - vecfx = _vec(f(_x)) - dx = imag(vecfx)/epsilon + vecfx_c = _vec(f(_x)) + dx = imag(vecfx_c)/epsilon Ji = _make_Ji( J, rows_index, cols_index, dx, colorvec, color_i, nrows, ncols) J = J + Ji diff --git a/test/out_of_place_tests.jl b/test/out_of_place_tests.jl index d3dca1b..ad32282 100644 --- a/test/out_of_place_tests.jl +++ b/test/out_of_place_tests.jl @@ -55,3 +55,17 @@ J = FiniteDiff.finite_difference_jacobian(f, x, Val{:central}, eltype(x)) J = FiniteDiff.finite_difference_jacobian(f, x, Val{:complex}, eltype(x)) @test J ≈ fill(1.0, 2, 1) @test J isa SMatrix{2,1} + +# The per-color closures used to assign variables that also live in the enclosing +# function's scope, which boxed the captures and made the return type `Any`. +# Static-array inputs stay uninferrable here: the `mapreduce(_, hcat, _)` accumulator +# grows one `SMatrix` column per color, so its type depends on `maximum(colorvec)`. +@testset "Type stability of the dense out-of-place jacobian" begin + g(x) = x .^ 2 .- 2 + x = [1.0, 2.0, 3.0] + @testset "$difftype" for difftype in (:forward, :central, :complex) + cache = FiniteDiff.JacobianCache(x, Val{difftype}, eltype(x)) + @test (@inferred FiniteDiff.finite_difference_jacobian(g, x, cache)) ≈ + Diagonal(2x) + end +end