Skip to content

Preserve inference for operator-backed linear caches - #1089

Draft
ChrisRackauckas-Claude wants to merge 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:codex/fix-operator-cache-inference
Draft

Preserve inference for operator-backed linear caches#1089
ChrisRackauckas-Claude wants to merge 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:codex/fix-operator-cache-inference

Conversation

@ChrisRackauckas-Claude

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

Copy link
Copy Markdown
Member

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

Depends on #1073 and #1087.

Summary

  • materialize and own convertible operator Jacobians for the default and concrete-matrix linear solver paths
  • preserve operator-backed caches for explicitly matrix-free linear solvers
  • refresh an owned concrete cache from the externally maintained operator without rebinding or aliasing its storage
  • unwrap AutoSpecialize problems only for cache consumers that actually use Enzyme-backed differentiation
  • add an end-to-end inference and solve regression for the default TrustRegion path, with Enzyme loaded before cache inference is checked
  • bump NonlinearSolveBase to 2.35.1

Root cause

OrdinaryDiffEq's reused W is represented as a MatrixOperator. Embedding the full operator-backed LinearCache in LinearSolveJLCache exceeded Julia's inference budget, so construct_linear_solver inferred an existential inner cache type even though the runtime value was concrete. This broke OrdinaryDiffEqBDF's existing @inferred cache-construction test after OrdinaryDiffEq commit 9a2f7e7b.

Default and factorization paths already require a convertible concrete matrix. Materializing an owned copy at initialization keeps their cache type inferable and prevents in-place factorization from modifying the externally owned operator storage. Explicit Krylov paths still retain the operator and remain matrix-free. When the ODE layer refreshes its operator, the LinearSolve extension now copies the new concrete coefficients into the owned cache buffer.

The minimum-compatibility CI graph also loads Enzyme before the operator regression, which selects AutoEnzyme as the default VJP backend. NonlinearSolveFirstOrder.__init previously unwrapped the problem eagerly for every AD-related cache. Recovering an AutoSpecialize function through its type-erased raw-function field then made the Jacobian, descent, and trust-region cache parameters abstract, even though the default direct-operator TrustRegion path does not use a VJP.

Problem unwrapping is now scoped to the cache that consumes each AD backend. Generic trust-region schemes unwrap only for Yuan and Bastin, which construct AD-backed operators; custom schemes retain the prior conservative behavior. Line-search, forcing, Jacobian, and Levenberg-Marquardt paths continue to unwrap when their selected backend actually requires Enzyme.

Validation

  • exact CI-floor graph on Julia 1.10.11 with the downgraded manifest and allow_reresolve=false: full NONLINEARSOLVE_TEST_GROUP=Core passed, including TrustRegion 1596/1596, LevenbergMarquardt 45/45, and SciMLOperator Jacobians 16/16
  • latest-compatible graph on Julia 1.12.6 after Pkg.update(): full NONLINEARSOLVE_TEST_GROUP=Core passed, including TrustRegion 1176/1176, LevenbergMarquardt 33/33, and SciMLOperator Jacobians 16/16
  • NONLINEARSOLVE_TEST_GROUP=Core julia +1.12 --project=lib/NonlinearSolveBase -e 'using Pkg; Pkg.test()' — passed
  • lib/NonlinearSolveFirstOrder/test/operator_jacobian.jl with the exact registered LinearSolve 4.3 dependency stack — passed (12/12, 2/2, 2/2)
  • OrdinaryDiffEqBDF test/inference_tests.jl with the exact registered LinearSolve 4.3 dependency stack — passed; the unmodified base reproduces the inference failure
  • repository-wide Runic --check . and git diff --check — passed

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

Copy link
Copy Markdown
Member Author

Investigation record for commit 20239c9c9:

  1. Reproduced the downstream OrdinaryDiffEqBDF inference failure on clean OrdinaryDiffEq aa895b563 with registered NonlinearSolveBase 2.35.0, NonlinearSolveFirstOrder 2.2.0, LinearSolve 4.3.0, and SciMLOperators 1.24.3.
  2. With that manifest frozen, OrdinaryDiffEq commit 3b799e583a93adc25b104718242778d3cf7364e1 passed and its child 9a2f7e7b79141951df24a0f2191a8aa20e94f49a failed. The child introduced reuse of the ODE W as an operator-backed inner Jacobian.
  3. Reduced the inference boundary to NonlinearSolveBase.construct_linear_solver: the runtime LinearSolveJLCache was concrete, while its inferred return type existentially abstracted the inner LinearCache only when it held the large MatrixOperator closure.
  4. Verified that materializing an owned matrix for default/concrete solvers restores exact inference. An initial implementation exposed the need to refresh existing concrete caches without assigning the operator into them; the final extension method preserves that ownership invariant.
  5. Ran the exact downstream regression, focused operator tests, both affected Core suites, and repository-wide Runic checks before publishing.

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

Copy link
Copy Markdown
Member Author

Pushed the Enzyme/minimum-compatibility follow-up in 27490c477.

The CI-only inference regression came from eager problem unwrapping after Enzyme selected AutoEnzyme as the default VJP backend. The direct SciMLOperator/default-TrustRegion path did not consume that VJP, but recovering the AutoSpecialize function through the type-erased raw-function field made the nested cache parameters abstract. Unwrapping is now scoped to the cache consumer; Yuan and Bastin retain the Enzyme-backed operator paths they require.

Local validation after the change:

  • Julia 1.10.11, exact downgraded CI-floor manifest, allow_reresolve=false, full NONLINEARSOLVE_TEST_GROUP=Core: passed; TrustRegion 1596/1596, LevenbergMarquardt 45/45, SciMLOperator Jacobians 16/16.
  • Julia 1.12.6, latest-compatible resolver after Pkg.update(), full NONLINEARSOLVE_TEST_GROUP=Core: passed; TrustRegion 1176/1176, LevenbergMarquardt 33/33, SciMLOperator Jacobians 16/16.
  • Repository-wide Runic --check . and git diff --check: passed.

The regression now loads Enzyme before the existing @inferred init(prob, TrustRegion()) assertion. No clean-base test failure was observed during this follow-up. The PR body now records its dependencies on #1073 and #1087.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Hosted downgrade classification for head 27490c4771a352791b8fdb42c204ea867c6a9e2a:

  • The downgrade action resolved the unstacked PR merge ref from root Project.toml, selecting the declared root floor LinearSolve 4.3.0.
  • Pkg.test(...; allow_reresolve=false) then incorporated the in-tree NonlinearSolveQuasiNewton 1.14.1, which requires LinearSolve 5, and stopped before loading tests with the expected 4.3.0 explicit-pin versus 5 source-package conflict (job 87761167898).
  • This PR explicitly depends on Restore root LinearSolve 5 floor #1087, which restores the root LinearSolve 5 floor. The exact stacked Julia 1.10.11 downgrade graph has already passed the full Core group locally, as recorded above and in the PR body.

No test, assertion, tolerance, dependency skip, or resolver bypass is being changed for this unstacked dependency failure.

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