Skip to content

Make the out-of-place jacobian's return type inferrable - #229

Merged
ChrisRackauckas merged 1 commit into
JuliaDiff:masterfrom
ChrisRackauckas-Claude:fix-oop-jacobian-inference
Aug 3, 2026
Merged

Make the out-of-place jacobian's return type inferrable#229
ChrisRackauckas merged 1 commit into
JuliaDiff:masterfrom
ChrisRackauckas-Claude:fix-oop-jacobian-inference

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Aug 3, 2026

Copy link
Copy Markdown

Please ignore until reviewed by @ChrisRackauckas.

Closes SciML/NonlinearSolve.jl#1092

Problem

The out-of-place finite_difference_jacobian infers to Any:

julia> g(x) = x .^ 2 .- 2
julia> x = [1.0, 2.0, 3.0];
julia> cache = FiniteDiff.JacobianCache(x, Val{:forward}, eltype(x));
julia> Base.return_types(FiniteDiff.finite_difference_jacobian, typeof.((g, x, cache)))
1-element Vector{Any}:
 Any

The calculate_Ji_forward / calculate_Ji_central / calculate_Ji_complex closures assign
variables that also exist in the enclosing function's scope (epsilon, vecfx, x_save,
dx, ...). An inner function assigning a name that exists in the enclosing scope binds the
outer variable, so those became captured variables with several assignment sites — which
closure conversion boxes. Reading a boxed capture gives Any, and that propagates through
mapreduce(calculate_Ji_*, hcat, ...) to the returned matrix. rows_index / cols_index
had the same problem via the comprehension that reassigns them.

@report_opt on the FiniteDiff path, before:

captured variable `epsilon` detected
captured variable `vecfx` detected
captured variable `rows_index` detected
captured variable `cols_index` detected
runtime dispatch detected: op::typeof(hcat)(%69::Any, %70::Any)::Any

Fix

Declare the closures' temporaries local so they stop aliasing the enclosing function's
variables, rename the few that had to stay distinct from a genuine capture, and give
vecfx / rows_index / cols_index a single assignment site each.

No behavioral change — the closures only ever wrote those outer variables incidentally, and
no branch read the values back afterwards.

After:

julia> Base.return_types(FiniteDiff.finite_difference_jacobian, typeof.((g, x, cache)))
1-element Vector{Any}:
 Matrix{Float64}

Static-array inputs stay uninferrable, which is a separate (and harder) problem: the hcat
accumulator gains one SMatrix column per color, so its type depends on the runtime value
of maximum(colorvec). Noted in a comment next to the new test.

Downstream motivation

This is the root cause of SciML/NonlinearSolve.jl#1092 — the
Jacobian cache is typed by this return value, so the whole solver cache inferred to an open
where-bound:

using NonlinearSolve, ADTypes, SciMLBase
prob = NonlinearProblem((u, p) -> u .^ 2 .- 2.0, [1.0])
@inferred SciMLBase.init(prob, TrustRegion(autodiff = AutoFiniteDiff()))   # ERROR before, OK after

Testing

Added an inference regression test to test/out_of_place_tests.jl covering all three
fdtypes. Ran the full GROUP=Core suite locally on Julia 1.10.11 — green:

Test Summary:    | Pass  Total  Time
Public API Tests |   15     15  0.2s
Test Summary:             | Pass  Broken  Total   Time
FiniteDiff Standard Tests |  219       1    220  35.6s
Test Summary:               | Pass  Total     Time
Color Differentiation Tests |   39     39  4m12.5s
Test Summary:      | Pass  Total   Time
Out of Place Tests |   28     28  40.1s
Test Summary:            | Pass  Total  Time
Cache Reuse Safety Tests |   27     27  4.5s

(the one Broken is pre-existing)

🤖 Generated with Claude Code

https://claude.ai/code/session_01MoeiVWqAEFJQK22kiXu48f

`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 <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MoeiVWqAEFJQK22kiXu48f
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Author

CI: test (Core, 1) is green. test (Downstream, 1) is red, but it is pre-existing on master — the last three master CI runs (29184179777, 29184168199, 29184127467) all show the same Downstream failure with Core passing.

The failure is unrelated to this diff:

LoadError: ArgumentError: array of type LinearAlgebra.Tridiagonal{Float64, Vector{Float64}} and size (10, 10)
  can not be filled with 1.0, since some of its entries are constrained.
  [1] fill! @ .../stdlib/v1.12/LinearAlgebra/src/special.jl:358 [inlined]
  [2] build_J_W(...) @ OrdinaryDiffEqDifferentiation .../src/derivative_utils.jl:1231

fill! on a structured matrix inside OrdinaryDiffEq's Rosenbrock alg_cache, nothing to do with finite_difference_jacobian.

@ChrisRackauckas
ChrisRackauckas marked this pull request as ready for review August 3, 2026 17:16
@ChrisRackauckas
ChrisRackauckas merged commit 9bae080 into JuliaDiff:master Aug 3, 2026
5 of 6 checks passed
ChrisRackauckas added a commit to SciML/NonlinearSolve.jl that referenced this pull request Aug 3, 2026
`inference_tests__item1.jl` only exercised in-place problems, so the
out-of-place instability in #1092 went unnoticed:

    prob = NonlinearProblem((u, p) -> u .^ 2 .- 2.0, [1.0])
    @inferred SciMLBase.init(prob, TrustRegion(autodiff = AutoFiniteDiff()))

inferred the cache's `jac_cache`/`descent_cache` to open `where`-bounds.
The root cause was FiniteDiff's out-of-place `finite_difference_jacobian`
returning `Any` (JuliaDiff/FiniteDiff.jl#229); the Jacobian cache is typed
by that return value, so the whole solver cache went with it.

Add the out-of-place problem to the existing testset and raise the
FiniteDiff lower bound to the release carrying the fix, so Downgrade CI
resolves a version where these tests hold.



Claude-Session: https://claude.ai/code/session_01MoeiVWqAEFJQK22kiXu48f

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

init(::NonlinearProblem, TrustRegion) not type-stable: Dogleg/JacobianCache infer to open where-bounds

2 participants