Skip to content

Add adaptive Jacobian reuse to first-order solvers#1072

Draft
ChrisRackauckas-Claude wants to merge 14 commits into
SciML:masterfrom
ChrisRackauckas-Claude:codex/jacobian-reuse-heuristics
Draft

Add adaptive Jacobian reuse to first-order solvers#1072
ChrisRackauckas-Claude wants to merge 14 commits into
SciML:masterfrom
ChrisRackauckas-Claude:codex/jacobian-reuse-heuristics

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Jul 15, 2026

Copy link
Copy Markdown
Member

Ignore this PR until it has been reviewed by @ChrisRackauckas.

Prompt

Add OrdinaryDiffEq-style Jacobian reuse heuristics to NonlinearSolve's first-order methods, benchmark standalone nonlinear-solver behavior before choosing a default, and keep OrdinaryDiffEq's separately owned Jacobian/iteration-matrix lifecycle unaffected.

Summary

  • add a public, documented JacobianReuse(max_age = 10, max_residual_ratio = 1) policy
  • keep exact Newton behavior as the default; opt in with jacobian_reuse = true or a configured policy
  • reuse a concrete factorization only while its Jacobian is unchanged
  • refresh after the age limit, residual deterioration or non-finite progress, a stale linear failure, a stale line-search failure, or a stale rejected descent/trust step
  • preserve a fresh Jacobian across repeated trust-region rejections at an unchanged state
  • reset policy state on reinit!; explicit step!(...; recompute_jacobian = true/false) decisions take precedence
  • expose the keyword on NewtonRaphson, TrustRegion, GaussNewton, LevenbergMarquardt, PseudoTransient, and the first-order polyalgorithms
  • bump NonlinearSolveFirstOrder to 2.3.0 and NonlinearSolve to 4.22.0

NewtonRaphson already constructs GeneralizedFirstOrderAlgorithm, so no second solver type is needed: the existing solver becomes adaptive only when the policy is supplied.

Prerequisite stack

The published branch is stacked on upstream master 60736b5b6, in this order:

  1. #1061: eliminate dense continuation-step allocations, cache the ForwardDiff destination, and avoid constructing full solution objects for every interior corrector
  2. #1073: use the public FunctionWrappersWrappers construction API and preserve inference
  3. #1057: preserve triangular Broyden initialization and the aliased matrix fallback
  4. the unchanged Jacobian-reuse feature commit
  5. #1087: restore the root LinearSolve 5 floor
  6. #1088: remove the duplicate BibTeX entry
  7. #1066: restore GPU Broyden inverse initialization

The feature stack has 12 commits. git range-diff marks every feature patch equal to the previously published stack; the new #1061 follow-up was moved beside its prerequisite. git diff --exit-code confirms that the reordered feature head is tree-identical to the locally tested pre-reorder tree. Two temporary test-only commits add failure-only allocation profiling for the unresolved Apple result and declare the Profile stdlib in the existing root test target. Every commit retains the required Chris Rackauckas co-author trailer.

Benchmark and default decision

BenchmarkTools medians, one BLAS thread, warmed solves:

Case Exact Newton Reuse result
dense diagonal, n=2 0.0089 ms 0.0096 ms (age 2)
dense diagonal, n=256 6.01 ms 2.65 ms (age 5)
trust region, n=256 7.01 ms 1.99 ms (age 5)
sparse Brusselator, 2048 unknowns 36.8 ms 26.6 ms (age 5)
23-problem Newton library 2.88 ms, 22/23 success 1.76 ms (age 2), same 22/23 success

For the sparse Brusselator, age 5 reduced Jacobian evaluations 4→3 and factorizations 3→2, while nonlinear steps increased 3→7. Policy sensitivity is real: the age-5 problem-library run regressed to 3.36 ms, and the n=2 case showed pure overhead. Matrix-free runs had identical Jacobian/factorization counts and inconsistent timing noise, so no matrix-free benefit is claimed.

Therefore exact Newton and the default polyalgorithm remain unchanged. jacobian_reuse = true selects a conservative age-10 bounded policy; users with expensive concrete Jacobians can tune age and residual ratio.

OrdinaryDiffEq isolation

OrdinaryDiffEqNonlinearSolve owns distinct J and W refresh decisions, including mass-matrix or γdt changes and convergence history across time steps. Its explicit recompute_jacobian Boolean overrides this standalone policy, while reinit! resets standalone reuse state.

The Robertson FBDF/TRBDF2 Jacobian-reuse test passed 8/8. Repeating both methods with NewtonRaphson(jacobian_reuse = true) remained accurate and successful; FBDF observed 10 Jacobian calls versus 115 W builds, and TRBDF2 observed 7 versus 129.

Apple ARM allocation follow-up

The first #1061 implementation depended on scalar replacement of temporary array wrappers. Caching the fixed-shape ForwardDiff destination removed that wrapper construction but did not resolve the combined Apple result: macOS prerelease still measured exactly 772.0 B/step for sweep and 400.0 B/step for arclength.

The cache-only follow-up factors generic cache iteration from solution construction. Sweep and arclength consume the reused cache directly for interior attempts; ordinary solve!, anchors, final landings, and reported failure originals still build complete solutions. Specialized polyalgorithm, no-init, and ForwardDiff caches keep their existing full-solution fallback.

Fresh Apple ARM prerelease CI on feature head 5b9e07da1 disproved the proposed root-cause attribution: the direct cache-only solve and other fast-path checks passed, but the end-to-end slopes remained exactly 772.0 and 400.0 B/step. The remaining allocations are therefore outside _solve_without_solution!. Current test-only head cfb8e2d18 profiles allocation types and NonlinearSolve stack sites only when the unchanged strict slope assertions fail; the follow-up declares Profile in the existing root [extras]/[targets] test environment. The completed stable-macOS profile identifies the dominant allocator upstream in LinearSolve's AppleAccelerateLUFactorization: repeated LU{Float64, Matrix, Vector{Int32}} construction contributed 25,984 bytes / 812 allocations in sweep and 11,648 bytes / 364 allocations in arclength, with additional tuple, Ref, range, and keyword-NamedTuple allocations. The profiled slopes were 957.28 and 496 bytes per step.

A focused upstream fix is now draft LinearSolve #1099. It reuses private factor, pivot, and status workspaces and adds direct macOS LTS/stable/pre allocation coverage. Its Apple-silicon CI now passes on Julia 1.10, 1.12, and 1.13: the refactorization-reuse suites passed 55/55, 59/59, and 59/59, while each mixed-precision suite passed 19 tests with its one pre-existing broken expectation. No cross-platform allocation fix is claimed in this downstream PR until #1099 is merged/released and the continuation profile is rerun against that dependency.

Validation

Feature-stack head: 5b9e07da11ab32181e86824ef0c33051f3989a29. Current test-only diagnostic head: cfb8e2d18e7fc2df972beab9c0bf5959acfa987d.

Final-tree local runs:

  • exact root allocation regression on Julia 1.12.6: cache-only solve 0 B and same-cache identity; sweep/arclength 0.0/0.0 B per step
  • NonlinearSolveBase Core on Julia 1.12.6 and 1.10.11: passed, including allocation fast paths 6/6; current also passed workspace 55/55, routing 12/12, operator 1/1, dampening 12/12, and dense LU 14/14
  • all 39 unchanged homotopy/arclength files: passed
  • NonlinearSolveBase QA: 18/18
  • ForwardDiff cache reinit regression: 1/1
  • root PolyAlgorithms from the preceding cached-destination tree: 107/107
  • Runic 1.7.0 --check --diff . and git diff --check: passed

The Jacobian-reuse feature patch itself is unchanged by range-diff. Its prior integrated validation passed adaptive reuse 64/64, General NLLS 472/472, NewtonRaphson 309/309, TrustRegion 1176/1176, PseudoTransient 51/51, LevenbergMarquardt 33/33, EisenstatWalker 27/27, sparse 12/12, mass-matrix 27/27, structured 16/16, and operator-Jacobian 14/14. The full final-tree FirstOrder Core run has since repeated these feature totals successfully. Fresh combined CI passed 78 checks before the Apple prerelease Core lane reproduced the unchanged 772.0/400.0 B/step failure. The exact diagnostic test passes locally on Julia 1.13.0-rc1 x86_64; the first diagnostic CI attempt exposed only the missing isolated-test declaration for Profile, now corrected at the current head. The completed profiling rerun isolated the remaining Apple allocations to LinearSolve's Accelerate LU wrapper construction; the downstream retest is pending LinearSolve #1099.

The canonical root Core run on this EPYC 7502 host stops before continuation at the separately audited Brown bad-Broyden @test_broken unexpected pass. The identical dependency/source graph fails on another EPYC model; this platform-sensitive clean-base issue remains tracked in #1056. No platform gate or weakened expectation is included.

No test threshold, tolerance, skip, broken-test designation, documentation failure mode, or unrelated public API was changed.

Copy link
Copy Markdown
Member Author

Implementation scratchpad for commit 6dd21578d:

  • preserve exact-Newton defaults and OrdinaryDiffEq's explicit J/W refresh ownership
  • add a stateless public policy plus per-cache residual/age tracking
  • refresh on bounded age, lost residual progress, or stale solve/globalization failures
  • preserve current Jacobians across unchanged-state trust-region rejections
  • forward the opt-in policy through first-order constructors and polyalgorithms
  • validate with Runic, Julia 1.10/1.12 targeted tests, FirstOrder Core/QA, root PolyAlgorithms, and FBDF/TRBDF2 integration
  • retain strict docs link checking; the local build reached document checks and then hit transient failures on pre-existing external links (timeout/HTTP 429)

Benchmarking showed wins for expensive dense/sparse factorizations but overhead and policy sensitivity on small/problem-library cases, so the solver default remains exact Newton.

@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the codex/jacobian-reuse-heuristics branch from 6dd2157 to 99348eb Compare July 15, 2026 13:57

Copy link
Copy Markdown
Member Author

Published final stacked head 99348eb0d.

Safety and history checks:

  • fetched upstream, fork, and origin immediately before push
  • upstream remained at 6264446f4
  • remote feature branch remained at the original single-author head 6dd21578d; no contributor commits intervened
  • git range-diff reports 6dd21578d = 99348eb0d for the feature patch
  • updated with an exact --force-with-lease against 6dd21578d
  • every stacked commit retains the required Chris Rackauckas co-author trailer
  • invalid Brown expectation commit is absent

Final local outputs include:

  • NonlinearSolveFirstOrder Core: passed (feature 64/64, Newton 309/309, TrustRegion 1176/1176, remaining groups passed)
  • NonlinearSolveFirstOrder QA: 18/18
  • Julia 1.10 focused feature suite after clean resolve: 64/64
  • root PolyAlgorithms: passed, including 46/46 problem-library assertions
  • root QA: 17/17
  • NonlinearSolveBase Core: passed
  • NonlinearSolveBase QA: 18/18
  • continuation inference: 26/26
  • Runic and git diff --check: passed

CI is now running on the combined stack, which includes both focused prerequisite fixes before the unchanged feature commit.

@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the codex/jacobian-reuse-heuristics branch from 99348eb to 63e043b Compare July 16, 2026 19:47
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Final-stack update at 63e043babea0cf86da82009c1b4c918a4ceeeb3e:

The PR description now records the actual stack, benchmark/default decision, baseline-failure investigations, and observed validation. Refreshed combined-stack CI is running.

Avoid temporary reshape wrappers for dense AutoForwardDiff Jacobians and same-buffer restructure calls, with direct warmed allocation regressions.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the codex/jacobian-reuse-heuristics branch from 63e043b to ed4d653 Compare July 17, 2026 00:23
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Published rebuilt combined head ed4d653e16a378505214b1f601028fda91ad90d7 with the Apple ARM allocation follow-up placed directly behind #1061.

Before replacing the old 63e043babea0cf86da82009c1b4c918a4ceeeb3e remote head, I fetched origin, upstream, and the fork, then audited the exact branch history plus PR commits, comments, and reviews. The PR still had the audited head, no reviews, and no comments from other contributors. The update used an exact --force-with-lease against that commit.

git range-diff marks every commit unchanged; only the new follow-up moved beside its prerequisite. git diff --exit-code confirmed that the reordered head is tree-identical to the locally tested tree. All 11 commits retain the required Chris Rackauckas co-author trailer.

New local evidence includes:

  • NonlinearSolveBase Core passed on Julia 1.12 and 1.10.11, including allocation fast paths 6/6
  • NonlinearSolveBase QA: 18/18
  • unchanged root continuation allocation regression: sweep/arclength 0.0/0.0 B per step
  • root PolyAlgorithms after the follow-up: 107/107
  • Runic 1.7.0 over the repository and git diff --check: passed

The superseded combined CI run finished 90 successful, 3 skipped, and 4 failed. Two Ubuntu failures were the already-audited Brown unexpected-pass behavior. The two macOS failures were the allocation issue addressed by the cached destination. Fresh CI is now running; Apple ARM must pass before the cross-platform fix is claimed.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the codex/jacobian-reuse-heuristics branch from ed4d653 to 5b9e07d Compare July 17, 2026 02:47

Copy link
Copy Markdown
Member Author

Published rebuilt combined head 5b9e07da11ab32181e86824ef0c33051f3989a29 with the cache-only corrector commit placed directly behind #1061.

Safety audit immediately before publication:

  • fetched origin, upstream, and the fork
  • upstream remained 60736b5b6
  • the remote feature branch remained exactly ed4d653e16a378505214b1f601028fda91ad90d7
  • PR Add adaptive Jacobian reuse to first-order solvers #1072 remained draft, with no reviews or comments from other contributors
  • updated only this branch using --force-with-lease pinned to that exact audited SHA

git range-diff marks all 12 patches equal; only the new allocation commit moved beside its prerequisite. The new head is tree-identical to the fully tested pre-reorder tree, and all commits have the required co-author trailer.

The prior Apple prerelease job is important negative evidence: it still measured exactly 772.0/400.0 B per sweep/arclength step after the cached Jacobian destination. The remaining allocation came from constructing a complete solution for every interior corrector. The new helper runs those correctors to completion in the reused cache; ordinary solve results, anchors, landings, and failure originals are unchanged.

Final local evidence includes cache-only 0 B with same-cache identity, end-to-end 0.0/0.0 B per step, Base Core on Julia 1.12 and 1.10, all 39 continuation files, Base QA 18/18, ForwardDiff cache reinit 1/1, Runic 1.7, and git diff --check. Fresh Apple ARM CI is required before claiming the cross-platform result.

Copy link
Copy Markdown
Member Author

Final-tree local integration update for 5b9e07da11ab32181e86824ef0c33051f3989a29:

  • GROUP=Core timeout 3600 julialauncher +1.12 --project=lib/NonlinearSolveFirstOrder -e 'using Pkg; Pkg.test()' passed.
  • Key totals observed in this run: adaptive reuse 64/64, general NLLS 472/472, NewtonRaphson 309/309, PseudoTransient 51/51, TrustRegion 1176/1176, Levenberg–Marquardt 33/33, Eisenstat–Walker 27/27, sparse autodiff 12/12, structured Jacobians 16/16, and SciMLOperator Jacobians 14/14.
  • The generated ignored lib/NonlinearSolveFirstOrder/Manifest.toml was removed; both worktrees are clean.

Fresh Apple ARM Core checks remain queued and are still required for the cross-architecture allocation claim.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Copy link
Copy Markdown
Member Author

Apple ARM prerelease result on feature head 5b9e07da1:

  • tests / Core (julia pre, macos-latest) failed at the unchanged strict continuation slopes: sweep 772.0 B/step, arclength 400.0 B/step.
  • The new direct cache-only solve/identity check and preceding continuation tests passed. This disproves the proposed attribution to per-corrector NonlinearSolution construction; the remaining bytes are outside _solve_without_solution!.
  • I have not changed the < 48 thresholds or silenced the test.
  • Test-only diagnostic commit ec549add5 adds Profile.Allocs aggregation by allocation type and first NonlinearSolve stack site only when a slope fails. It passed locally, including profiler smoke output, and is now triggering a fresh Apple run.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Copy link
Copy Markdown
Member Author

Diagnostic rerun correction at cfb8e2d18:

  • The first diagnostic CI attempt did not reach profiling on Julia 1.12 because Pkg.test() isolates undeclared test dependencies and Profile was absent from the root test target.
  • Profile is now declared in the existing root [extras] and targets.test; no test/Project.toml was created.
  • The exact allocation test passes locally on Julia 1.13.0-rc1 (exit 0), and TOML parsing confirms both declarations.
  • The unchanged strict < 48 B/step assertions remain in place. A fresh Apple ARM run is now queued to produce the allocation-type/site profile.

Copy link
Copy Markdown
Member Author

Apple allocation profile root cause and upstream follow-up:

  • Stable macOS measured sweep/arclength slopes of 957.28 / 496 B per step with the unchanged < 48 assertions.
  • LU{Float64, Matrix, Vector{Int32}} construction in the Apple Accelerate LinearSolve path contributed 25,984 bytes / 812 allocations for sweep and 11,648 bytes / 364 allocations for arclength. The profile also showed repeated tuple, Ref, range, and keyword-NamedTuple construction.
  • This maps to AppleAccelerateLUFactorization rebuilding its LU wrapper and LAPACK workspace representation on each fresh refactorization, not to the NonlinearSolve cache-only result path.
  • Draft upstream fix: SciML/LinearSolve.jl#1099. It keeps factors, pivots, and LAPACK status in a private mutable cache and adds macOS LTS/stable/pre allocation tests.
  • No threshold was changed or suppressed. The cross-platform downstream claim remains pending LinearSolve macOS CI and a NonlinearSolve rerun against the fixed dependency.

Copy link
Copy Markdown
Member Author

Upstream Apple-silicon validation update for LinearSolve #1099:

  • Julia 1.10: refactorization reuse 55/55; mixed precision 19 passed with 1 pre-existing broken expectation
  • Julia 1.12: refactorization reuse 59/59; mixed precision 19 passed with 1 pre-existing broken expectation
  • Julia 1.13: refactorization reuse 59/59; mixed precision 19 passed with 1 pre-existing broken expectation

These are direct LinearSolve macOS tests of the fixed Apple Accelerate path. They do not replace the required NonlinearSolve end-to-end continuation rerun. The combined PR therefore still makes no downstream cross-platform allocation claim until #1099 is merged/released and this branch is retested against that dependency.

Copy link
Copy Markdown
Member Author

Fresh exact-graph validation of the stacked #1057 prerequisite:

  • Baseline NonlinearSolve 60736b5b6 against LinearSolve 0496f135 on Julia 1.10.11 reproduced root Broyden at 94 passed / 1 failed / 1 errored / 19 broken: Generalized Rosenbrock alg Make iip a compile-time constant #4 returned residual 4.3999999999999995, and Brown alg Make iip a compile-time constant #4 unexpectedly passed its existing broken expectation.
  • Applying the Preserve triangular solves for Broyden Jacobian inversion #1057 patch unchanged made NonlinearSolveBase Core pass (including workspace 53/53), NonlinearSolveQuasiNewton Core pass (Broyden 810/810, Klement 297/297, limited-memory Broyden 135/135, dedicated Rosenbrock regression 2/2), and changed root Broyden to 95 passed / 0 failed / 1 errored / 19 broken.
  • Generalized Rosenbrock therefore passes with the exact Preserve triangular solves for Broyden Jacobian inversion #1057 code already included in this combined stack. The only remaining root error is the separately tracked platform-sensitive Brown unexpected pass in #1056.
  • Runic and git diff --check passed; temporary validation patches and generated files were removed.

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.

2 participants