Skip to content

nlls-gram 2.7.0: unified solver contracts, shape-aware AD defaults, float32/GPU correctness - #30

Merged
jlperla merged 22 commits into
mainfrom
refactor/unify-solver-contracts
Jul 26, 2026
Merged

nlls-gram 2.7.0: unified solver contracts, shape-aware AD defaults, float32/GPU correctness#30
jlperla merged 22 commits into
mainfrom
refactor/unify-solver-contracts

Conversation

@jlperla

@jlperla jlperla commented Jul 26, 2026

Copy link
Copy Markdown
Member

Release branch for 2.7.0. Breaking, and the breaks are deliberate.

Carried instances (the headline refactor)

Metric and preconditioner instances are pytrees carried in LMState, and all
adaptive policy moved into the per-step callback. LMSolveAction is now
LMAction.

Shape-aware implicit-AD defaults, restored

2.6 had collapsed both of 2.4's AD defaults onto Cholesky. Each loss was a
real defect:

  • square: 2.4 factored J directly; Cholesky computes the same map
    through B'B, at cond^2. Measured in float64: 2.6e-10 at cond 1e4, 7.1e-6
    at 1e6, and 1.5e-1 at 1e8, where the right answer is 1.3e-9 and the
    solve still reports CONVERGED. The new LU config restores it.
  • rectangular: 2.4 resolved to SVD. The undamped dual is singular
    whenever the small side is rank deficient, which padded zero residuals
    produce by construction. spooky's growth_recursive_advanced.py returned an
    all-NaN Jacobian under the Cholesky default and returns the analytically
    correct one under SVD. test_float64_svd_ad_solver_near_duplicate_rows
    documents that exact pathology but pins ad_solver=SVD() explicitly, so it
    kept passing while the default moved out from under it.

auto now resolves square -> LU, rectangular -> SVD, which is 2.4's rule.

float32 on GPU

XLA:GPU serves float32 dot_general from TF32 tensor cores on Ampere: a
10-bit mantissa, ~1e-3. Forming a Gram matrix already squares the condition
number, so TF32 spends ~3 decimal digits before the factorization. Eleven
tests failed on an RTX 3090 and passed on CPU; the same eleven failed before
this branch, so it was long-standing, not a regression.

Every product the package owns is now pinned to Precision.HIGHEST. That is
necessary but not sufficient: the Jacobian is differentiated through the
caller's residual matmuls, so callers must also set
jax_default_matmul_precision="highest" (3.3e-4 -> 1.6e-7 on a dense float32
solve). Documented in the tuning guide.

Removed

min_damping/max_damping. Both guarded non-failures -- unbounded damping
ends at MAX_STEPS with the best iterate, and _converged already blocks a
stall reading as convergence via xtol_met & info.accepted. The internal
anti-underflow floor at finfo(dtype).tiny stays: the update is
multiplicative, so zero is absorbing.

Verification

CPU GPU (RTX 3090)
nlls_gram 215 passed 227 passed, 0 failed (was 11 failed)
tinydiffeq 251 passed 266 passed, 0 failed
spooky replication 385/385 leaves bit-identical --float32 clean, first-ever baseline
kernels implicit AD 1.4e-10 neutrality 9.5e-11

New float32 module (dtype purity + accuracy across the config matrix), an
x64-with-float32-inputs promotion test over 18 configs, and the first tests
for Nystrom/Woodbury/ShermanMorrison/Padded preconditioners.

Downstream

Publishing this breaks published tinydiffeq 2.2.0, which declares
nlls-gram>=2.4.0 and passes 17 kwargs that no longer exist. v2.6.0 still
accepts them, so the break starts here. A tinydiffeq release with a
>=2.7.0 floor follows immediately.

jlperla and others added 22 commits July 25, 2026 00:46
@jlperla @claude
refactor: shared LM substrate -- split utility.py, merge state types,…
d1efc92 
… base class

utility.py (1099 lines) splits into lm_types.py (status/hyperparams/state/
info/action/context/result), solve_loop.py (the jitted while_loop driver, its
Python mirror, save_steps buffers), multi_start.py (MultiStart, DrawNNXModule,
the three drivers), and utilities.py (tree selection/masking, static-key
hashing, residual canonicalization).

LMState/RidgeLMState and LMInfo/RidgeLMInfo merge into one pair each: the
ridge fields default to None and stay None for the metric solver, where a
None subtree costs nothing in the while_loop carry -- the same design the
optional resid/Jt/aux slots already relied on.

New LevenbergMarquardtBase holds what both solvers did identically: solve()
and its custom_jvp implicit-AD wrapper, _solve_impl/_multi_start_impl, the
callback-action plumbing, dense Jacobian assembly, and value-based __eq__/
__hash__. RidgeLevenbergMarquardt subclasses it and keeps only the ridge
contract, through five hooks (_validate_tolerances, _solve_lm_state,
_initial_ad_point, _check_action_state, _apply_action_state). The metric
solver ports in a later commit; here it only follows the moved imports.

Pure refactor: 596 passed, 12 skipped, unchanged from baseline.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor: linear solvers own their cache and their solve
b4c0c7f 
solver_config.py becomes linear_solvers.py and grows the contract the configs
were missing. Each config now supplies:

  new_cache(m, n, n_m, dtype) -> the reject-step cache pytree, or None
  prepare(Subproblem)         -> StepSolver(grad, velocity, solve, accel_rhs,
                                            make_cache)

so init() no longer branches on a solver-name string to allocate a cache, and
update() no longer carries three parallel if/elif blocks that each redefine
solve_step/accel_rhs. The QR path's extra machinery -- its backward-stable Q2
velocity route and the corrected semi-normal refinement -- lives behind
StepSolver.velocity instead of a fourth closure the other branches lacked.
CholeskyCache/QRCache move alongside the configs that build them, and
_resolved_solver() (the string the branches keyed on) is gone.

Dispatch still happens once at trace time, so the compiled program is
unchanged; Subproblem pins metric-callback outputs to the residual dtype in
one place rather than at eleven call sites.

Constructor validation drops the isinstance checks on user-supplied types and
the reserved metric_factory stub, keeping the mathematical invariants (ridge >
0, the damping schedule, a CG stopping rule) and the size checks. The ridge
solver still rejects a None ridge in update() -- with the merged LMState that
is now a legal state for the metric solver, so it is a real contract check
rather than defensive noise.

ridge_lm.py: 2043 -> 1114 lines. 596 passed, 12 skipped, unchanged.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor: metrics and preconditioners own their adaptive state
70cfcda 
Both hook base classes gain the same optional pair:

  prepare(theta, ctx) -> traced state pytree, or None (the default)
  rebuild(ctx)        -> traced predicate gating a rebuild, True by default

The output rides on lm_state (metric_state/precond) and comes back as
ctx.metric_state / ctx.preconditioner_state, rebuilt on accepted steps and
reused across rejected ones. Whether a hook is stateful is a static property
of its class, so the slots and their lax.cond compile away entirely for the
stateless default.

BlockEigenPreconditioner stops reading ctx.args[args_key]: it holds
blocks_fn + permutation and eigendecomposes in prepare, so its state is no
longer threaded through the residual args and rebuilt by hand from a solve
callback. rebuild() is the knob for declining a refresh -- a stale
preconditioner only changes the CG iteration path, never the converged step.

MetricContext becomes SolverContext and moves to lm_types: it now serves the
metric, the preconditioner, AND the linear solver, and carries the two new
state slots. Under implicit AD the hooks are prepared once at the returned
solution and not differentiated, the documented freeze contract.

594 passed, 12 skipped (four BlockEigenPreconditioner tests for the old
args-threading contract replaced by two for the new one).

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor!: one Metric, one solver menu, one implicit-AD contract
3fefd37 
The metric solver moves onto everything the ridge solver already had.

GramMetric is gone. With M = F'F its four optional callables are the ridge
Metric's factor ops at identical cost -- inv_sqrt IS factor_solve, "whiten"
and "unwhiten" ARE the factor solves -- so both solvers now take one Metric,
and _validate_metric_requirements (the which-callbacks-does-this-solver-need
matrix) has nothing left to check. metric_from_cholesky/metric_from_diagonal
become CholeskyMetric/DiagonalMetric, and repeated_shifted_dense_metric
becomes RepeatedFactorMetric.from_gram, with the old zero-pad tail now the
solver's free block weighted by Metric.free_scale.

gram_lm.py becomes metric_lm.py and subclasses LevenbergMarquardtBase. The
eight-value string menu becomes the same typed configs the ridge solver takes:
Cholesky(form=auto|gram|normal), QR(), CG(precond), GramCG(precond). QR is now
the damping-row form, which is rank-safe and so subsumes both the old qr and
augmented_qr. The five preconditioner keywords (dual_, normal_, whitened_,
ad_solver_, and the factory) collapse into one typed Preconditioner whose
space the config names, which deletes the mutual-exclusion matrix and the
three missing_*_preconditioner branches along with it.

Seven AD tangent implementations become four. The undamped AD operator is
singular on whichever side the problem is deficient in, so each Krylov rule is
now offered only where its operator is invertible and says so loudly
otherwise -- previously CG returned a silently wrong tangent for n > m.
regularized_normal_cg becomes CG(precond, penalty=...), so ad_solver_penalty
stops being a free-floating kwarg that errors with five of seven methods.
Reverse mode through the unpenalized normal rule uses the push-through
identity N^+ = B'(BB')^{+2}B, since a cotangent does not lie in range(B').

Deleted: lsmr.py, recycled_cg.py, WhitenedPreconditioner, and the metric-LM
constructor's other 17 keywords (30 -> 13). quasiseparable and the state-space
metric move to experimental/ as StateSpaceMetric. The dual-space preconditioner
helpers become classes (Nystrom/Woodbury/ShermanMorrison/Padded).

Tests: 4697-line test_gram_lm.py plus test_normal_solvers, test_lsmr,
test_recycled_cg, test_augmented_qr, test_implicit_geometry,
test_metric_factory, test_preconditioner_factory, test_ad_solver_methods and
test_jacobian_mode (9700 lines) replaced by a 400-line test_metric_lm.py where
every check is against a closed form or an independent reference, plus a
trimmed float64 suite. 149 passed, 13 skipped in 48s (was 596/195s).

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: value-hashable continuation callback, and guards for the whole c…
d48195f 
…lass of bug

The solve loop marks the callback a jit STATIC argument, so
ridge_continuation() returning a fresh closure meant every construction
recompiled the entire loop. It now returns a frozen RidgeContinuation
dataclass, value-hashable on (ridge_floor, decrease, grad_rtol, stall_rtol),
so equal schedules share one compiled loop.

tests/test_compilation.py pins the whole class of regression, measured
against the jitted loop's own compilation cache rather than a proxy. It
asserts one compilation for: solvers rebuilt repeatedly with equal configs,
traced values and loop controls changing, a reused metric across solver
rebuilds, a freshly constructed IdentityPreconditioner, a rebuilt continuation
callback, and the realistic driver that reconstructs the solver every
iteration. It also asserts the negative direction -- a changed shape, linear
solver, static scalar, or continuation schedule IS a different program -- so
the guards cannot pass by never compiling anything.

Also confirmed while writing these: the +1 residual evaluation per solve is
init()'s eager sizing pass, not a retrace.

161 passed, 13 skipped.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
docs: restructure around the two solvers; update benchmarks; 2.7.0
bd41d25 
Docs go 4252 -> 1743 lines. index.md leads with the one thing that
distinguishes the solvers (where the selection of the root lives) as a table,
then the residual interface and the solver menu; the eight per-solver math
subsections it used to carry duplicated tuning_guide and utilities.

gauss_newton.md and metrics.md become metric_lm.md (the damping geometry and
its minimum-norm limit, with the rank-deficiency table that says which
ad_solver is valid at which shape). utilities.md -- a grab-bag of eleven
unrelated helpers, a third of them for deleted features -- becomes metrics.md,
covering just the two hook types and why one must be exact and the other need
not be. Every mkdocstrings stub moves into api.md, so a removed name breaks
one file instead of six. tuning_guide.md and implicit_ad.md are rewritten
around what survives; llms.txt (orphaned, not in the nav) is deleted.

Benchmarks move to the typed configs and the new metric constructors; the
metric-factory case becomes a prepare()-based metric, which is what it was
measuring. The augmented_qr benchmark goes with the feature. 37 CPU
benchmarks run clean.

Also deleted: FAILED_IMPLICIT_AD_PLAN.md and the stale benchmarks/results
artifacts from the 2026-07 investigation.

mkdocs build --strict passes. 161 passed, 13 skipped.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
chore: rename the metric benchmark, drop stale name references
49376ab 
Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: reject solver configs in roles they cannot fill, and three AD bugs
1ae43b2 
From an external review. Five real bugs, all reproduced before fixing:

1. The ridge solver accepted GramCG, whose dual operator never sees the
   penalty rows -- it returned the UNPENALIZED step (1.0 where the ridge step
   is 0.5) with no complaint. It only rejected Cholesky(form="gram").
2. The normal-CG implicit VJP's push-through transpose reused the
   parameter-space preconditioner for its two dual solves, handing an
   m-vector to a hook expecting n. Those solves are now unpreconditioned,
   which is what residual space requires here.
3. ad_solver=QR() was accepted and silently ran dense Cholesky; the ridge
   solver treated QR/SVD/GramCG as normal CG; forward SVD() failed only later
   inside update.
4. Both init methods built the hook context WITHOUT lm_state, so a metric or
   preconditioner prepare() reading ctx.lm_state.damping -- documented as
   available -- crashed on the first call.
5. BlockEigenPreconditioner unconditionally read ctx.lm_state.ridge, which is
   None under LevenbergMarquardt, so it could not serve the metric solver at
   all.

(3) is fixed structurally rather than case by case: each config declares
supports_forward / supports_ad / supports_penalty, and one shared
_validate_configuration in the base rejects the combination at construction.
Cholesky.form and jacobian_mode now reject unknown values instead of falling
through to a different algorithm, and the metrics validate free_scale > 0 and
size >= 0 (a non-positive free_scale makes the whitening noninvertible).

Also from the review: _cold_state and _block_sizes were still duplicated
between the solvers and move to the base; LMState.recycle and
LevenbergMarquardt.ridge were dead after their features were removed; the
ridge module and LMInfo docstrings still described blockdiag(F, I) and
metric-LM grad_norm as ||J'r||, both stale since whitening became shared.

Tests: the reject-step cache test was silently SKIPPING (its fixture no
longer produced a rejected step) and now exercises the path; the padded-SVD
test checked only the primal, where ad_solver affects only the tangent; added
coverage for role rejection, free_scale's effect on step and tangent, and a
tall normal-CG VJP with a dimension-specific preconditioner (which is what
catches bug 2). 165 passed, 12 skipped.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: stale hook state after a callback move, and four AD/dtype defects
09a1ae9 
From a second (JAX-focused) review.

The serious one: _apply_action invalidated the Jacobian cache when a callback
replaced x or args, but not metric_valid or precond_valid. update() sets
metric_valid = ~improved, so after a REJECTED step it is True; a callback that
then installed a new x left the metric state prepared at the old iterate in
place and the next step reused it. The metric defines the subproblem, so that
silently solved a different problem in a different geometry. All three
validity flags now get the same treatment.

Also fixed:
- solve() called hyperparams() with no dtype, so CG(tol=None) resolved against
  the JAX default float. On a float32 problem under enabled x64 that gave
  tol=1e-10 -- unreachable, since float32 eps is 1.2e-7 -- and every inner CG
  burned its full maxiter for the whole solve. The sentinel now resolves in
  _cast_hyper, where the residual dtype is known.
- The solve callback was not passed through _hashable_hook, so a callback
  written to the frozen-dataclass pattern the docs recommend died with a raw
  "Non-hashable static arguments" error if it held an array. draw and accept
  already had this.
- ad_solver=None under CG(penalty=...) inherited the preconditioner but not
  the penalty, silently discarding the option -- and leaving the AD solve on
  the singular operator whose explicit form the guard rejects.
- Cholesky's docstring called gram-vs-normal "a cost choice, not a semantics
  choice". True in exact arithmetic; false in floating point for m > n, where
  the dual carries m-n structural zero eigenvalues and loses ~1e-2 relative at
  damping=1e-14 against the normal form's 1e-15. form="auto" never picks gram
  there, so the default was always safe; the docstring now says which is.

Documented rather than changed: only result.x/aux/p carry tangents. info
(including info.loss), lm_state, steps, and the histories get exact zero,
which jax.grad reports as a silent zero on an otherwise natural bilevel
objective. Damping and step norms are path artifacts, not properties of the
root, so the contract is right -- it just was not written down.

Dead: canonicalize_ad_preconditioner (unreferenced, documents the removed
string API). Metric.norm's "must match to floating-point accuracy" contract
was vacuous -- neither solver calls it.

test_compilation.py was blind to the multi-start drivers' own jit caches; it
now counts all three and adds the trap the file was missing: a residual lambda
rebuilt per call recompiles every time (3 solves, 3 programs), while a
module-level one compiles once. Also pins save_steps making max_steps static,
and a stateful preconditioner reused across solver rebuilds.

169 passed, 12 skipped.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor: remove RepeatedFactorMetric.from_gram; callers factor the s…
55891a8 
…hifted Gram themselves

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
docs: drop 'nugget' wording from the quasiseparable docstring
ec330ec 
Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
docs: pytree instance-state design (reviewed; decisions settled)
29c1e72 
Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor!: metric and preconditioner instances are pytrees carried in…
c8e11e2 
… LMState

Instances register via register_pytree_dataclass (bypass-unflatten, statics
as type-tagged aux), ride in lm_state.metric/preconditioner, and every
traced read goes through the carried instance -- rebuild is calling the
constructor again inside the solve callback. Compile identity keys on
instance structure, so equal-config fresh instances share one loop. Deletes
prepare/rebuild, hook-state fields and machinery, block_eigen_state, and
the ridge_continuation factory (AnnealRidge, with init_state, replaces it);
renames LMSolveAction/LMSolveContext to LMAction/LMContext. A ridge-solver
metric change suppresses that step's convergence test and stales the solver
caches; the metric solver only stales caches; preconditioner refreshes are
free. ad_solver=CG(None, ...) inherits the carried forward preconditioner
at the solution with pinned AD knobs; failed lanes differentiate at the
pre-loop instances.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
test+docs: pin the carried-instance contract; rewrite adaptation docs
cb12131 
New pins: same-treedef different-valued instances give different answers
under the shared compile; callback swaps neither recompile nor (for
preconditioners) invalidate; ridge metric swaps suppress that step's
convergence while metric-solver swaps do not; the untouched-instance
identity short-circuit emits no comparison ops; GramCG equation-space
preconditioning with a mid-solve refresh plus a CG-budget schedule; the
implicit tangent uses the carried metric at the solution.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: reject nested unregistered instances at solver construction
2154dca 
Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: freeze solvers after construction; weak-type check in the instan…
0643a49 
…ce guard

A mutated solver attribute would keep the stale static key and silently
reuse another configuration's compiled loop, so assignment now raises. The
callback structure guard compares leaf weak types too: a weak/strong
scalar swap passes the carry's physical checks while retracing the body
under different promotion rules.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
refactor!: drop min_damping/max_damping; keep the anti-underflow floor
294593b 
Both knobs guarded failure modes that are not failures. Without an upper cap,
repeated rejections can overflow damping to inf in float32 -- but the step
then goes to zero, every trial is rejected, x stays at the last accepted
iterate, and the solve ends at MAX_STEPS with the best point it found. The
dangerous case, a stall reported as convergence, is already blocked in
_converged by `xtol_met & info.accepted`: a stalled solver only produces
rejected steps, so xtol cannot fire. max_damping turned MAX_STEPS into
MAX_STEPS.

min_damping as a user knob was equally inert. It could only ever RAISE the
floor (the resolver takes a max against the dtype floor), so the documented
advice to "lower it (e.g. 1e-12)" was impossible to follow and would have
introduced the endgame truncation it warned about. Damping falling freely is
the point: the endgame wants it to vanish so the step approaches Gauss-Newton
and the minimum-norm limit. Every test that pinned min_damping=1e-12 passes
without it, against a floor 296 orders of magnitude lower.

What remains is the internal floor at finfo(dtype).tiny, which guards
something genuinely irreversible: the update is multiplicative, so a damping
in a backend's flush-to-zero range is absorbing -- 0 * damping_increase stays
0 and the solver could never re-damp again. That floor is far below the scale
at which damping still perturbs the Gram diagonal, so it is an anti-underflow
backstop, not regularization; RidgeLevenbergMarquardt is the tool for that.
A callback clamp replaces both knobs for the rare problem needing a bound,
which is where 2.7.0 puts adaptive policy anyway.

Behavior is unchanged for every caller that did not pass them: min_damping
None already resolved to tiny and max_damping None already skipped the upper
clamp. spooky's full replication is bit-identical across the change (387
leaves; only train_time moves) and a warm repeat run still writes zero
.jax_cache entries.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix!: restore shape-aware implicit-AD defaults (LU when square, SVD o…
1688a9d 
…therwise)

2.6/2.7 collapsed the AD rules onto Cholesky: ad_solver=None resolved to the
assembled normal/dual factorization at every dense shape. That lost both of
2.4.0's shape-appropriate defaults, and each loss is a real defect.

Square (a DAE root, a determined system): 2.4.0 factored J itself via
ad_solver="direct". Cholesky computes the same map through B'B, at
cond(B)^2. Measured on a square constraint in float64, the tangent error
tracked cond^2 * eps instead of cond * eps -- 2.6e-10 at cond=1e4, 7.1e-06 at
1e6, and 1.5e-01 at 1e8, where the correct answer is 1.3e-09 and the solve
still reports CONVERGED. The new LU config restores it. Squareness is the
exact condition that makes a plain solve valid: the tangent is unique, so no
norm is minimized and the metric selects nothing, which is also why the square
path can skip the whitening round-trip the rectangular rules need. One
factorization serves both directions. LU is ad_solver-only (the damped
forward subproblem is SPD at every shape) and raises on a rectangular system
rather than guessing.

Rectangular: 2.4.0 resolved to SVD, and this is the more serious loss. The
undamped dual is singular whenever the small side is rank deficient, which is
not exotic -- test_float64_svd_ad_solver_near_duplicate_rows already documents
the growth-model pathology where a converged simulation duplicates its
late-horizon states to ~1e-13. That test pinned the behavior with an explicit
ad_solver=SVD(), so it kept passing while the default silently went to
Cholesky. spooky's notebooks/growth_recursive_advanced.py, which
differentiates a neural solve against a closed-form steady state, returned an
all-NaN Jacobian under the Cholesky default and returns the analytically
correct one under SVD.

Verified across consumers: tinydiffeq 251 passed with the DAE reverse-mode
regression closed (vjp-vector16-dae +18.4% -> -1.6%, worst DAE case +22.4% ->
+3.2%); spooky's replication bit-identical on 385/385 non-timing leaves with
zero new .jax_cache entries; kernels' multicountry AD experiment unchanged at
its analytic neutrality bound (1.40e-10 both ways); spooky's advanced AD
example NaN -> correct.

Cholesky() remains available as the opt-in for a rectangular system whose
small side is known to have full rank, trading tangent accuracy for speed.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
test: float32 coverage for solvers, implicit AD, and preconditioners
cd471a9 
nlls-gram has no ad/metric solve-dtype knob -- the 2.6/2.7 refactor removed
linear_solve_dtype, metric_solve_dtype and ad_dtype -- so a float32 problem
must stay float32 end to end, at float32 accuracy, with no promotion to
recover conditioning. Nothing pinned that.

tests/test_float32.py runs at default precision and covers, for each config,
both dtype purity and accuracy against a float64 reference, so a silently
promoted solve and a silently wrong one each fail: dense forward solvers
(Cholesky auto/gram/normal, QR), matrix-free (CG, GramCG), the implicit-AD
rules on a rectangular system (auto->SVD, SVD, Cholesky, GramCG) and on a
square one (auto->LU, LU, SVD), the metrics, and every preconditioner.

test_float64_subprocess.py gains the harder direction: x64 ENABLED with
float32 inputs, where Python scalars default to f64 and can promote the whole
compute silently. Across 18 configurations no primal, jvp or vjp jaxpr
contains an f64 compute op; f64 appears only in explicit convert_element_type
at the call boundary.

This also gives NystromPreconditioner, WoodburyPreconditioner,
ShermanMorrisonPreconditioner and PaddedPreconditioner their first tests --
only BlockEigen and Identity had any. The apply contract is checked for
float32, finiteness and SPD (v'Mv > 0 plus symmetry on a random pair), and
GramCG is run end to end under each dual preconditioner to confirm a
preconditioner changes the CG path but not the converged step.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
fix: pin internal matmuls to HIGHEST so float32 is not served from TF32
5a58142 
XLA:GPU serves float32 dot_general from TF32 tensor cores by default on
Ampere and later: a 10-bit mantissa, so ~1e-3 relative where float32 is
~1e-7. Forming a Gram or normal matrix already squares the condition number,
so paying TF32 on top spends about three decimal digits before the
factorization starts. That is not a trade this package should be making
silently on a user's behalf.

It was doing exactly that. Eleven tests fail on an RTX 3090 and pass on CPU --
cross-solver step agreement, the metric factor round-trip, the block-eigen
apply against its dense inverse, and the CG-vs-Cholesky implicit tangent --
every one of them a comparison routing through a product. They are not a
regression: the same eleven fail at 0643a49, before any of this branch. They
are also not tolerance problems. On the 3090:

    default (TF32)                        11 failed, 20 passed
    JAX_DEFAULT_MATMUL_PRECISION=highest  31 passed

Every product the package owns now goes through utilities.mm (or passes
precision=HIGHEST to einsum): the Gram and normal assembly, the matrix-free CG
operator, the QR back-substitution, the metric factor applications, all six
preconditioners, the state-space metric scans, and the implicit-AD tangent
solves. Unconditional, with no opt-out knob -- a float32 solve now answers the
same on GPU as on CPU, where the setting is a no-op. The metric round-trip is
the tell for why this is the right layer: factor_apply is a matmul and
factor_solve is a triangular solve, so each op looked fine alone while the
round-trip drifted.

tests/conftest.py sets jax_default_matmul_precision globally, which the
package's own pinning does NOT cover: the tests build dense references of
their own, and without it those would stay TF32 and the comparison would fail
on the reference side. To keep that from masking the library change, a new
test traces a solve whose residual contains no matmul -- so every dot_general
in the jaxpr is the solver's -- and asserts each carries HIGHEST. Asserting on
the jaxpr rather than on numbers makes it device-independent, so CPU CI guards
the GPU path; there is no GPU runner.

The float32 tolerances are untouched. They sit above the TF32 floor, which is
why all 26 float32 tests passed on GPU while the library was degraded -- they
now pass for the right reason.

docs/tuning_guide.md gains a section on the trap, what the package pins, what
it cannot reach (matmuls in the caller's residual), and the failure signature:
CPU and GPU agreeing to about three digits and no further, worst in the
tangent.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
docs: the GPU matmul setting is required, not advisory
791fafa 
Measured on the 3090: the package pinning its own products leaves a dense
float32 solve at 3.3e-4 relative, because the Jacobian is differentiated
through the caller's residual matmuls and arrives already carrying TF32
error. Setting jax_default_matmul_precision=highest takes the same solve to
1.6e-7. Since a residual here is usually a network or a kernel evaluation,
that is the normal case.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla @claude
test: assert the carried-instance contract directly, not through conv…
11246e2 
…ergence

test_callback_refresh_reaches_the_next_inner_solve failed in CI on x86 while
passing on ARM. It is not a regression -- it fails identically at 0643a49 --
and it was never visible before because the test is branch-only (added by
c8e11e2) and CI runs only on main and PRs, so this branch's 21 commits had
never been through it.

The mechanism is fine. Swapping a preconditioner into the carried state
produces a BIT-IDENTICAL step to a solver built with that instance, on both
x86 and ARM (measured: ||swapped - exact|| = 0.0).

What the old assertion actually measured was the float32 endgame. The residual
is linear, so with a starved maxiter=3 budget the solve reaches the float32
loss floor, no trial step strictly improves the loss any more, every step is
rejected, and damping ratchets to ~3e4. Whether gtol=1e-5 sits above or below
that floor is platform arithmetic: ARM reached 4.5e-6, x86 stopped at 6.0e-5
and stayed there at 60, 200 and 600 steps. Exact-from-start converges to ~1e-6
on both, which is what made the difference look like a mechanism failure
rather than a stalled endgame.

So the convergence assertion is replaced by the contract it was standing in
for: the step taken with a swapped-in instance equals the step a solver built
with it takes from the same state, and differs materially from the stale one
so the check is not vacuous. The end-to-end callback path still asserts the
refreshed instance is the one carried out, that a starved scrambled run does
not converge, and that refreshing strictly improves stationarity -- which
holds on both platforms without depending on where the float32 floor lands.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014og23CSBQfdHGNCfA21F8x
@jlperla
jlperla merged commit c10aa1e into main Jul 26, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jlperla