Preserve inference for operator-backed linear caches - #1089
Preserve inference for operator-backed linear caches#1089ChrisRackauckas-Claude wants to merge 2 commits into
Conversation
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
|
Investigation record for commit
|
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
|
Pushed the Enzyme/minimum-compatibility follow-up in The CI-only inference regression came from eager problem unwrapping after Enzyme selected Local validation after the change:
The regression now loads Enzyme before the existing |
|
Hosted downgrade classification for head
No test, assertion, tolerance, dependency skip, or resolver bypass is being changed for this unstacked dependency failure. |
Ignore this PR until it has been reviewed by @ChrisRackauckas.
Depends on #1073 and #1087.
Summary
AutoSpecializeproblems only for cache consumers that actually use Enzyme-backed differentiationTrustRegionpath, with Enzyme loaded before cache inference is checkedRoot cause
OrdinaryDiffEq's reused
Wis represented as aMatrixOperator. Embedding the full operator-backedLinearCacheinLinearSolveJLCacheexceeded Julia's inference budget, soconstruct_linear_solverinferred an existential inner cache type even though the runtime value was concrete. This broke OrdinaryDiffEqBDF's existing@inferredcache-construction test after OrdinaryDiffEq commit9a2f7e7b.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
AutoEnzymeas the default VJP backend.NonlinearSolveFirstOrder.__initpreviously unwrapped the problem eagerly for every AD-related cache. Recovering anAutoSpecializefunction through its type-erased raw-function field then made the Jacobian, descent, and trust-region cache parameters abstract, even though the default direct-operatorTrustRegionpath 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
allow_reresolve=false: fullNONLINEARSOLVE_TEST_GROUP=Corepassed, includingTrustRegion1596/1596,LevenbergMarquardt45/45, andSciMLOperator Jacobians16/16Pkg.update(): fullNONLINEARSOLVE_TEST_GROUP=Corepassed, includingTrustRegion1176/1176,LevenbergMarquardt33/33, andSciMLOperator Jacobians16/16NONLINEARSOLVE_TEST_GROUP=Core julia +1.12 --project=lib/NonlinearSolveBase -e 'using Pkg; Pkg.test()'— passedlib/NonlinearSolveFirstOrder/test/operator_jacobian.jlwith the exact registered LinearSolve 4.3 dependency stack — passed (12/12, 2/2, 2/2)test/inference_tests.jlwith the exact registered LinearSolve 4.3 dependency stack — passed; the unmodified base reproduces the inference failureRunic --check .andgit diff --check— passed