Pick a GPU-compatible default AD when u0 lacks fast scalar indexing

[ChrisRackauckas-Claude]

What changed and why

Solving a JLArray-backed NonlinearProblem under JLArrays.allowscalar(false) fails
because NonlinearSolveBase's default AD/Jacobian selection
(select_forward_mode_autodiff/select_reverse_mode_autodiff/select_jacobian_autodiff)
always tries ForwardDiff-based backends first, and ForwardDiff's dual-seeding
scalar-indexes the state. When ArrayInterface.fast_scalar_indexing(u0) is false,
candidate backends are now filtered to finite-differencing ones
(AutoFiniteDiff/AutoFiniteDifferences); if none is loaded, the ArgumentError names
the fix. NonlinearSolveBaseForwardDiffExt's additional_incompatible_backend_check
applies the same rule for an explicitly-requested AutoForwardDiff/AutoPolyesterForwardDiff.

Round 2 (responding to a Devin SWE-2 first-pass review, round1/REVIEW.md): the reviewer
raised one blocking correctness concern (could not execute code) and one blocking CI
failure. Both are now resolved — see below.

Review finding 1 (matrix-free/Krylov VJP path) — investigated, NOT confirmed as a bug

The reviewer argued (by static reading only, no execution available) that because
vjp_autodiff resolves to AutoFiniteDiff() for non-scalar-indexable state, any
matrix-free path (NewtonRaphson(linsolve = KrylovJL_GMRES()), or TrustRegion with
radius_update_scheme = Yuan/Bastin, which eagerly compute J'fu via a VJP operator)
would crash during JacobianOperator's eager prepare_vjp, because
DifferentiationInterface's generic PullbackSlow fallback builds Jacobian columns via
basis(x, j), which does a scalar setindex!.

I ran all four configurations the reviewer named, on JLArray under allowscalar(false):
NewtonRaphson(linsolve=KrylovJL_GMRES()), TrustRegion(linsolve=KrylovJL_GMRES()),
TrustRegion(radius_update_scheme=Yuan), TrustRegion(radius_update_scheme=Bastin)
(the last two also combined with Krylov). All pass and return the correct root, e.g.:

=== NewtonRaphson(linsolve=KrylovJL_GMRES()) ===
retcode=Success u=[2.0000000000000036, 2.0]
=== TrustRegion(radius_update_scheme=Yuan) dense ===
retcode=Success u=[2.0000000000000027, 2.0]
=== TrustRegion(radius_update_scheme=Bastin) dense ===
retcode=Success u=[2.0, 2.0]

Root cause of the discrepancy with the reviewer's static analysis: DifferentiationInterface
ships ext/DifferentiationInterfaceGPUArraysCoreExt, which overrides
DI.basis(a::AbstractGPUArray, i) to wrap the single-element seed write in
GPUArraysCore.@allowscalar:

function DI.basis(a::AbstractGPUArray{T}, i) where {T}
    b = similar(a); fill!(b, zero(T))
    @allowscalar b[i] = oneunit(T)
    return b
end

This extension auto-loads whenever GPUArraysCore is in the session (it is, transitively
via JLArrays/GPUArrays), so the PullbackSlow per-column pushforward-based pullback
is already GPU-safe — the reviewer's read of basis (and of _prepare_pullback_aux,
which calls the zero-argument basis(x) that never indexes at all) didn't account for
this extension, understandably, since they had no shell to load it and check.

I still added the four configurations above as regression tests in
misc_tests__item11.jl (8 new @tests) to lock in this currently-correct behavior,
per the reviewer's "add tests" request.

Review finding 2 (downgrade CI) — confirmed, root-caused, fixed

Reproduced exactly: pinning JLArrays=0.3.0, ArrayInterface=7.19.0, GPUArrays=11.5.5, GPUArraysCore=0.2.0, LinearSolve=5.6.0, FiniteDiff=2.32.1, DifferentiationInterface=0.7.3, ADTypes=1.13.0, ForwardDiff=1.0.1, KernelAbstractions=0.9.44, Adapt=4.5.2 (the exact
versions the downgrade-sublibraries job installed, read from its log) while keeping
NonlinearSolveBase/NonlinearSolveFirstOrder/SciMLJacobianOperators on local dev
paths (this PR's actual code) reproduces the CI failure byte-for-byte:

ERROR: LoadError: DimensionMismatch: matrix is not square: dimensions are (4, 2)
 [2] cholesky!(A::JLArray{Float64, 2}, ::LinearAlgebra.NoPivot; check::Bool)
 ...
 [7] _init_default_cacheval(alg::LinearSolve.DefaultLinearSolver, ...)
     @ LinearSolve .../src/default.jl
 ...
 [17] init(..., alg::MoreTrustRegionDescent{...}, J::JLArray{Float64,2}, ...)
     @ NonlinearSolveBase .../src/descent/more_trust_region.jl:1062

Root cause: TrustRegion's default descent (MoreTrustRegionDescent) solves its
subproblem through an augmented, non-square [J; √λ·D] system (4×2 here) via
LinearSolve.DefaultLinearSolver. At LinearSolve 5.6.0–5.13.0, the default-algorithm
infrastructure (_init_default_cacheval, a @generated function) eagerly builds the
cache for every algorithm in DefaultAlgorithmChoice, including CholeskyFactorization,
before any algorithm is actually chosen — and that eager build calls cholesky!(A)
unconditionally, which requires a square matrix. For ordinary Array/Matrix, a
specialized dispatch path avoids ever touching this generic fallback for non-square input;
for JLArray in this old LinearSolve range, it does not. Not a JLArrays version
issue: I bisected by varying only LinearSolve while holding every other package
(including JLArrays) at the downgrade floor — bumping JLArrays alone to the newest
0.3.x (0.3.3) did not fix it; LinearSolve 5.10.0–5.13.0 still fail, LinearSolve
5.14.0 and above (tested up to current 5.18.2) pass.

Fix: raise NonlinearSolveFirstOrder's LinearSolve compat floor from "5.6" to
"5.14" (the verified minimum). Bumped the package's patch version (2.10.0 → 2.10.1)
for the compat-only change. Re-ran the exact faithful downgrade reproduction above with
this floor applied (LinearSolve pinned to 5.14.0, everything else unchanged at its
downgrade floor) and the full, now-expanded misc_tests__item11.jl (14 assertions) passes
cleanly.

I did not bump NonlinearSolveBase's own (lower, "5.4") LinearSolve floor:
MoreTrustRegionDescent/construct_linear_solver live there, so the same crash could in
principle occur there too, but NonlinearSolveBase has no JLArray test to exercise it
and its own downgrade job doesn't fail, so I left it alone rather than widen the diff —
flagging it here for a maintainer to decide.

Vector-unchanged evidence (round 1, re-confirmed)

solve on a Vector NonlinearProblem with NewtonRaphson()/TrustRegion() — output
captured on master and on this branch, byte-identical:

GeneralizedFirstOrderAlgorithm retcode=Success u=[2.000000000000002, 2.0] resid=[8.881784197001252e-15, 0.0] stats=SciMLBase.NLStats(5, 6, 5, 5, 5)
GeneralizedFirstOrderAlgorithm retcode=Success u=[2.000000000000002, 2.0] resid=[8.881784197001252e-15, 0.0] stats=SciMLBase.NLStats(5, 6, 6, 5, 5)

Failing-before / passing-after (round 1, re-confirmed)

Before (src reverted to master, keeping only the test file):
ERROR: LoadError: Scalar indexing is disallowed from ForwardDiff.seed! on the
JLArray. After (fix applied): passes.

Test-group tails (Julia 1.10.12, this branch, round 2)

• NonlinearSolveFirstOrder GROUP=Core (full suite, ~9000 assertions across all
  algorithms, now including the 14-assertion JLArray test): all pass, e.g.

  JLArray default AD without scalar indexing |   14     14  23.2s
  Native bounded methods                     | 5989   5989  6m40.2s
  NewtonRaphson                               |  435    435  8m16.3s
  TrustRegion                                 | 1596   1596  14m31.2s
  LevenbergMarquardt                          |   45     45  56.6s
  SciMLOperator Jacobians                     |   14     14  11.0s
  

  No Error/Fail/Some tests did not pass anywhere in the run.
• NonlinearSolveBase GROUP=Core: all pass, no Error/Fail entries (e.g. Bounds transform (#955) 138/138, construct_extension_jac scalar and matrix u0 10/10, Dense LU refactorization allocations 14/14).
• NonlinearSolveFirstOrder GROUP=QA:

  Test Summary: | Pass  Total   Time
  QA            |   18     18  34.1s
  

• NonlinearSolveBase GROUP=QA: still blocked locally by the pre-existing, unrelated
  SciMLJacobianOperators registry-lag issue noted in round 1 (confirmed reproduces with
  this PR's diff reverted).
• Runic on all changed files: no diff (already formatted). typos: clean.

Non-blocking review items

Reviewed all; left as-is per the reviewer's own "minor polish" / "user-explicit backends
are allowed to fail" characterization — in particular, AutoSparse(AutoForwardDiff())
passed explicitly by a user on a non-scalar-indexable state is not intercepted by
additional_incompatible_backend_check/is_finite_differences_backend (neither unwraps
AutoSparse). Fixing it would touch the same selection logic again for a case the user
opted into explicitly; flagging it as a follow-up rather than widening this PR.

What I did not verify

• Real GPU hardware (CUDA/Metal/ROCm) — only JLArrays.JLArray (CPU-backed, GPU-array
  semantics).
• NonlinearSolveBase GROUP=QA locally (pre-existing registry-lag gap, round 1).
• Whether NonlinearSolveBase's own lower LinearSolve floor (5.4) hits the same
  MoreTrustRegionDescent/JLArray crash — plausible by the same root cause, but
  untested since no local test exercises it.
• Downstream sublibraries and the All/Everything groups.

Anything a reviewer should push back on

• The LinearSolve floor bump ("5.6" → "5.14") is a compat-only change verified by
  bisection on this machine, not by reading LinearSolve's own changelog/PR history for
  the actual fixing commit — if a maintainer knows the exact PR, a tighter/more precisely
  justified floor might be possible.
• Whether NonlinearSolveBase's LinearSolve floor should also move, per the note above.

Please ignore until reviewed by @ChrisRackauckas.

Risk assessment

• Risk: low
• Blast radius: NonlinearSolveBase's default AD/Jacobian selection and
  NonlinearSolveFirstOrder's LinearSolve compat floor. No public API signature
  changes; behavior changes only when ArrayInterface.fast_scalar_indexing(u0) == false
  (GPU-style arrays) or when installed with a very old (pre-5.14) LinearSolve; Vector
  output confirmed byte-identical.
• Evidence: failing-before/passing-after test output; full GROUP=Core pass for both
  packages; byte-identical Vector-output comparison; the downgrade-CI failure faithfully
  reproduced locally and confirmed fixed by the LinearSolve floor bump alone (bisected
  against JLArrays separately to rule it out as the cause).
• Independent review: pending
• Merge: needs human review — touches default AD selection and a dependency compat
  floor on a widely-used code path

🤖 Generated with Claude Code (model: claude-sonnet-5)

https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

[ChrisRackauckas]

🤖 Automated comment from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris.

Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): CHANGES, risk medium.

Full review

VERDICT: CHANGES
RISK: medium

PR: #1368 (head ba7a6fd, merge-base 0055e31). No linked/closing issue, so there was no parent or tracking issue to check sibling work against.

Blocking findings

1. The JVP path for non-scalar-indexable state has regressed from exact ForwardDiff to finite differences, and an explicit jvp_autodiff = AutoForwardDiff() is now overridden. (Confirmed by running.)
   • Where: lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl:143-149 and lib/NonlinearSolveBase/src/autodiff.jl:79 / :185-191. additional_incompatible_backend_check and _autodiff_candidates feed every selector: Jacobian, forward-mode (JVP), and reverse-mode (VJP). Only dense Jacobian construction scalar-indexes, because ForwardDiff.seed! does duals[idx] = …. ForwardDiff's pushforward/JVP on a JLArray is GPU-safe and exact.
   • Evidence (scratch/probe2.jl, JLArrays.allowscalar(false), run against both master and the PR):

     # both envs
     pushforward AutoForwardDiff()   maxerr=0.0
     pushforward AutoFiniteDiff()    maxerr=7.55e-8
     # master
     NR Krylov default AD                Success nsteps=5 resid=1.79e-6 jvp_ad=AutoForwardDiff
     NR Krylov jvp=ForwardDiff explicit  Success nsteps=5 resid=1.79e-6 jvp_ad=AutoForwardDiff
     # PR
     NR Krylov default AD                Success nsteps=5 resid=4.53e-6 jvp_ad=AutoFiniteDiff
     NR Krylov jvp=ForwardDiff explicit  Success nsteps=5 resid=4.53e-6 jvp_ad=AutoFiniteDiff
         WARNING: The chosen AD backend `AutoForwardDiff()` does not support the chosen problem ... detected `AutoFiniteDiff()`
     

     The problem was Float32, n=200, a coupled A*u + u.^3 - b. scratch/probe_master_rev.log also shows that NewtonRaphson(linsolve = KrylovJL_GMRES()) on a JLArray already succeeds on master. That makes matrix-free Newton–Krylov, the main large-GPU configuration, a working path that this PR silently degrades. The warning text is also misleading: it says the backend package may not be loaded.
   • Rule: AGENTS.md says "Honor … JVP/VJP callbacks and AD backends" and "Krylov paths must work through operators without materializing the Jacobian". The PR overrides a backend the user chose explicitly and that works, and only the operator path needed no change at all.
   • Required change: apply the non-scalar-indexing exclusion only where a Jacobian is actually materialized. One way is the selection used for dense construct_jacobian_cache. Another is to materialize through a pushforward-per-column path, which is GPU-safe because DI's basis is @allowscalar-wrapped. Do not apply it in select_forward_mode_autodiff / additional_incompatible_backend_check for JVPs. When no concrete Jacobian is formed (Krylov / concrete_jac = false), the default JVP for a JLArray must stay AutoForwardDiff, and an explicit jvp_autodiff = AutoForwardDiff() must not be replaced. Add tests that assert the selected backend on the operator path (for example sol.alg.jvp_autodiff isa AutoForwardDiff for NR+GMRES on a JLArray) and that the explicit choice produces no warning. Today's tests only check retcode and root, so they cannot detect this degradation.

Non-blocking findings

• The VJP filter is broader than needed. (Confirmed by running, scratch/probe3.jl.) For VJP selection _autodiff_candidates discards everything except FD. On a JLArray, DI.pullback with AutoZygote is exact (maxerr 0.0). AutoEnzyme (both modes) fails with EnzymeRuntimeActivityError. Master's default reverse pick, Enzyme, would therefore have failed, so the PR is an improvement there. But a loaded Zygote is skipped in favour of FD. Consider filtering by "known to scalar-index" (ForwardDiff/Polyester, plus Enzyme if desired) rather than "keep only FD", and say so in the comment at autodiff.jl:183-184.
• The _nonscalar_state_hint error path is untested, and the selectors have no unit test in NonlinearSolveBase (the package that changed). All new tests live in NonlinearSolveFirstOrder and go through solve.
• The NR+Krylov test is not a regression test. misc_tests__item11.jl:20-22 also passes on master (see probe_master_rev.log), so it guards nothing this PR changed. The TR+Krylov, Yuan and Bastin cases do need the Jacobian selection change.
• AutoSparse note in the PR body is moot. NewtonRaphson(autodiff = AutoSparse(AutoForwardDiff())) errors on both master and the PR with "Specifying a sparse AD type for Nonlinear Problems was removed in v4", independent of this change.
• TrackedArray state is also caught. ArrayInterface.fast_scalar_indexing is also false for Tracker.TrackedArray / ReverseDiff.TrackedArray (ArrayInterface exts), so a tracked u0 that reaches selection would also switch to FD. In-repo solve_up overloads intercept tracked inputs first, so this is suspected harmless; not exercised.
• Compat and version mismatch. NonlinearSolveBase src behaviour changed without a version bump, while NonlinearSolveFirstOrder was bumped to 2.10.1. Its new test needs the new Base behaviour, but its NonlinearSolveBase = "2.54" compat allows the old Base. It is fine in-repo via [sources]. The repo does not usually bump per PR (for example 44e6a9a4), so this is minor.
• Two concerns in one PR. The LinearSolve floor raise ("5.6" → "5.14") is a compat change bundled with the AD-selection change. It is justified because the new TrustRegion/JLArray test needs it, and I reproduced the bisection endpoint (see "What I ran").
• PR body attribution does not follow CLAUDE.md. The body has no first-line UNREVIEWED banner, and the footer lacks the harness version and is not in the prescribed format. The commit trailers are correct.
• Unrelated red CI jobs. CI on the head has red jobs: Core (lts, ubuntu/macos) solve!() inference, Trim (macos), and SCCNonlinearSolve Core isconcretetype(eltype(original)). The same jobs are red on master at the merge-base, in runs https://github.com/SciML/NonlinearSolve.jl/actions/runs/37942671576 and https://github.com/SciML/NonlinearSolve.jl/actions/runs/37266049518. They are not attributable to this PR.

What I ran

All on Julia 1.12.7 in scratch environments that Pkg.develop the PR's lib/ (or a git archive of origin/master lib/). The checkout itself was not modified.

• New test file, fail before / pass after. I ran misc_tests__item11.jl in a @testset. PR: item11 | 14 14 pass. Master: errors with Scalar indexing is disallowed at item11.jl:8. This reproduces the author's failing-before / passing-after claim.
• Selection and solver probes (scratch/probe.jl). These ran on the PR, the PR with Zygote+Enzyme loaded, and master with both loaded.
  • Vector selection is unchanged: Jacobian/forward → AutoForwardDiff, explicit ForwardDiff is kept.
  • On a JLArray the PR picks AutoFiniteDiff for all three selectors.
  • With the PR, JLArray solves succeed for iip NR, TR and NR+GMRES, oop LevenbergMarquardt, PseudoTransient, and explicit AutoForwardDiff (which falls back to FD). On master the iip NR and TR solves, LevenbergMarquardt, PseudoTransient and explicit AutoForwardDiff all fail with scalar indexing; iip NR+GMRES already succeeds there. The AutoSparse case errors on both (see Non-blocking).
  • A coupled n=40 Float64 problem was compared against a Vector reference for NR, TR, NR+GMRES and TR-Yuan. The PR agrees to 1.4e-12, 7e-15, 3.4e-12 and 1.4e-12 respectively, with residuals ≤ 7e-12. On master only the Krylov case runs; the other three throw scalar indexing.
• scratch/probe2.jl: ForwardDiff vs FD pushforward on a JLArray, and Float32 n=200 NR+GMRES with the default and with explicit jvp_autodiff, on master and the PR. This is blocking finding 1.
• scratch/probe3.jl: DI pullback/pushforward on a JLArray with Zygote, Enzyme Reverse/Forward, and FiniteDiff. Zygote is exact; Enzyme fails with runtime activity. FiniteDiff called directly through DI with no preparation fails with a MethodError on DI's PushforwardFast path. The solver uses prepared operators and works anyway, as the Krylov solves show.
• LinearSolve floor claim. PR Base + SciMLJacobianOperators + master FirstOrder (floor 5.6), with LinearSolve pinned:
  • 5.13.0: DimensionMismatch: matrix is not square: dimensions are (4, 2), 3 pass / 1 error.
  • 5.14.0: 14/14 pass.
  • JLArrays was 0.3.3 and SciMLBase 3.57.0 in both. This confirms the 5.13 → 5.14 boundary with the other dependencies current. I did not reproduce the full downgrade-floor set.
• CI. I read the head's checks with gh pr checks and gh run view --log-failed, and compared them with master runs.
• Cleanup. No Julia processes were left running (checked with pgrep).

What I did not verify

• Real GPU hardware (CUDA/Metal). Only JLArrays was used.
• The author's full GROUP=Core / GROUP=QA runs for NonlinearSolveBase and NonlinearSolveFirstOrder. I did not run them, per the host limits. CI's FirstOrder and downgrade jobs on the head are green.
• The exact downgrade-sublibraries floor set (only LinearSolve was pinned).
• SimpleNonlinearSolve, NonlinearSolveQuasiNewton, arclength and wrappers. These also call select_jacobian_autodiff and get the same behaviour change, but I did not exercise them on a JLArray.
• The _nonscalar_state_hint error message, i.e. the case where no FD backend is loaded.

---

🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium
Conversation: local Claude session 3c311569-dda6-4687-bc9f-65febb212c33 (fleet master)