From 76e60d3dfaf968e7aa14c6e847500f939efc751f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 23 Sep 2026 15:25:51 +0200 Subject: [PATCH 1/3] fix(profiling): accept the library's new solver/stats kwargs in the certified solver injection (PyAutoArray#566) Co-Authored-By: Claude Fable 5.1 --- .../library_solver_injection.py | 36 ++++++++ scripts/misc/test/test_fixed_light_library.py | 90 +++++++++++++++++++ 2 files changed, 126 insertions(+) diff --git a/scripts/misc/likelihood_breakdown/library_solver_injection.py b/scripts/misc/likelihood_breakdown/library_solver_injection.py index ec704161e..ff0063eef 100644 --- a/scripts/misc/likelihood_breakdown/library_solver_injection.py +++ b/scripts/misc/likelihood_breakdown/library_solver_injection.py @@ -110,6 +110,22 @@ An ordered callback is a host round trip and serialises the batch, so it is a **diagnostic pass**, never a timed one: the caller runs the reporting pass once and then times and traces with ``report=None``. + +``solver`` and ``stats`` (PyAutoArray #566) +------------------------------------------- + +PyAutoArray #566 gave the library's entry point ``solver="pdip"`` (which +positive-only solver to run) and ``stats=None`` (a dict it reports solver +diagnostics into). The wrapper declares both, so :func:`_assert_signature_covers` +passes, and forwards them only to a library that accepts them +(:func:`_library_kwargs`) — the harness runs against the pre-#566 library too. + +On the NumPy path both are forwarded as given. On the JAX path the injected +certified route *replaces* the library's solver, so ``solver`` does not reach +the injected solve; and the route's PDIP **fallback** is defined as the +library's PDIP, so it calls the library with ``solver="pdip"`` explicitly — a +library config selecting its own certified solver must not change what route +(d) falls back to — and ``stats=None``. """ from __future__ import annotations @@ -167,6 +183,16 @@ def _assert_signature_covers(original, patched) -> None: ) +def _library_kwargs(original, **values) -> dict: + """The subset of ``values`` whose names ``original`` declares as parameters. + + Used to forward the #566 ``solver``/``stats`` keywords only to a library + that has them, so the same wrapper calls the old and the new signature. + """ + params = inspect.signature(original).parameters + return {name: value for name, value in values.items() if name in params} + + def certified_reconstruction_from( data_vector, curvature_reg_matrix, @@ -224,6 +250,10 @@ def certified_reconstruction_from( certified = jnp.any(out["certified"]) + # Route (d)'s fallback is the library's PDIP by definition: pin it, so a + # library config selecting another positive-only solver cannot change it. + fallback_kwargs = _library_kwargs(original, solver="pdip", stats=None) + def _library_pdip(): return original( data_vector=data_vector, @@ -232,6 +262,7 @@ def _library_pdip(): xp=xp, fingerprint=fingerprint, factor=factor, + **fallback_kwargs, ) return jax.lax.cond(certified, lambda: x_active, _library_pdip) @@ -273,6 +304,8 @@ def certified_solver_injected( original = inversion_util.reconstruction_positive_only_from counts = {"jax": 0, "numpy": 0} + # Read once: which of the #566 keywords the installed library accepts. + forwarded = frozenset(_library_kwargs(original, solver=None, stats=None)) def patched( data_vector, @@ -281,6 +314,8 @@ def patched( xp=np, fingerprint=None, factor=None, + solver="pdip", + stats=None, ): if not xp.__name__.startswith("jax"): counts["numpy"] += 1 @@ -291,6 +326,7 @@ def patched( xp=xp, fingerprint=fingerprint, factor=factor, + **{k: v for k, v in (("solver", solver), ("stats", stats)) if k in forwarded}, ) counts["jax"] += 1 return certified_reconstruction_from( diff --git a/scripts/misc/test/test_fixed_light_library.py b/scripts/misc/test/test_fixed_light_library.py index dcb01a7c3..de5216567 100644 --- a/scripts/misc/test/test_fixed_light_library.py +++ b/scripts/misc/test/test_fixed_light_library.py @@ -29,6 +29,7 @@ from __future__ import annotations +import inspect import re import subprocess import sys as _sys @@ -192,6 +193,95 @@ def test_numpy_path_is_delegated_untouched(): assert np.allclose(np.asarray(got, dtype=float), np.asarray(expected, dtype=float)) +def _stub_new(calls): + def stub( + data_vector, + curvature_reg_matrix, + settings=None, + xp=np, + fingerprint=None, + factor=None, + solver="pdip", + stats=None, + ): + calls.append({"solver": solver, "stats": stats}) + return data_vector + + return stub + + +def _stub_old(calls): + def stub( + data_vector, curvature_reg_matrix, settings=None, xp=np, fingerprint=None, factor=None + ): + calls.append({}) + return data_vector + + return stub + + +@pytest.mark.parametrize("new_signature", [True, False]) +def test_injection_installs_against_the_old_and_new_library_signature(monkeypatch, new_signature): + """PyAutoArray #566 added ``solver``/``stats``; the harness must run on either side of it.""" + from autoarray.inversion.inversion import inversion_util + + calls = [] + stub = (_stub_new if new_signature else _stub_old)(calls) + monkeypatch.setattr(inversion_util, "reconstruction_positive_only_from", stub) + + q = np.ones(4) + stats = {} + with lsi.certified_solver_injected(3) as counts: + inversion_util.reconstruction_positive_only_from( + data_vector=q, curvature_reg_matrix=np.eye(4), xp=np, solver="certified", stats=stats + ) + assert counts == {"jax": 0, "numpy": 1} + if new_signature: + assert calls == [{"solver": "certified", "stats": stats}] + else: + assert calls == [{}], "the old library must not be handed solver/stats" + assert inversion_util.reconstruction_positive_only_from is stub + + +def test_fallback_pins_the_library_pdip_solver(): + """Route (d)'s fallback is the library's PDIP, whatever solver the library config selects.""" + original = _library_positive_only() + assert "solver" in inspect.signature(original).parameters, ( + "this test needs the PyAutoArray #566 library on the path (source activate.sh)" + ) + seen = [] + + def spy( + data_vector, + curvature_reg_matrix, + settings=None, + xp=np, + fingerprint=None, + factor=None, + solver="", + stats="", + ): + seen.append((solver, stats)) + return original( + data_vector=data_vector, + curvature_reg_matrix=curvature_reg_matrix, + settings=settings, + xp=xp, + fingerprint=fingerprint, + factor=factor, + solver=solver, + stats=stats, + ) + + Q, q = _constrained_qp(60, seed=0) + jax.jit( + lambda Q_, q_: lsi.certified_reconstruction_from( + q_, Q_, n_passes=1, fallback=True, original=spy, xp=jnp + ) + )(jnp.asarray(Q), jnp.asarray(q)) + assert seen and all(s == ("pdip", None) for s in seen) + + # --------------------------------------------------------------------------- # 2. The two library facts the cell's routes stand on # --------------------------------------------------------------------------- From 0473c6adca685c3d650ee7f84c2f792dcdbf7238 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 23 Sep 2026 15:28:45 +0200 Subject: [PATCH 2/3] fix(profiling): forward the library's solver/stats kwargs in the numpy solver injection too (PyAutoArray#566) Co-Authored-By: Claude Fable 5.1 --- .../fixed_light_numpy_solvers.py | 25 +++++-- .../test/test_fixed_light_numpy_solvers.py | 66 +++++++++++++++++++ 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/scripts/misc/likelihood_breakdown/fixed_light_numpy_solvers.py b/scripts/misc/likelihood_breakdown/fixed_light_numpy_solvers.py index 307392b58..1e5dd4394 100644 --- a/scripts/misc/likelihood_breakdown/fixed_light_numpy_solvers.py +++ b/scripts/misc/likelihood_breakdown/fixed_light_numpy_solvers.py @@ -611,6 +611,18 @@ def numpy_solver_injected(solver, *, label: str): # works against a pre-lever-3 library too — where the parameter does not exist # and forwarding it would be a `TypeError`. original_takes_factor = "factor" in inspect.signature(original).parameters + # Likewise PyAutoArray #566's `solver` / `stats`: declared by the wrapper (so the + # signature still covers the library's) and forwarded on the delegated JAX path + # only to a library that has them. This is `library_solver_injection._library_kwargs` + # inlined: that module imports JAX, which this module must never do (see + # `test_the_numpy_modules_import_no_jax_at_any_level`). On the numpy path the + # injected kernel *is* the solve, so the library's `solver` choice does not reach + # it. The wrapper's keyword must be spelled `solver` (the library passes it by + # name), so the injected kernel is held as `injected`. + original_forwarded = frozenset( + name for name in ("solver", "stats") if name in inspect.signature(original).parameters + ) + injected = solver def patched( data_vector, @@ -619,12 +631,17 @@ def patched( xp=np, fingerprint=None, factor=None, + solver="pdip", + stats=None, ): if xp.__name__.startswith("jax"): counts["jax"] += 1 kwargs = {} if original_takes_factor: kwargs["factor"] = factor + for name, value in (("solver", solver), ("stats", stats)): + if name in original_forwarded: + kwargs[name] = value return original( data_vector=data_vector, curvature_reg_matrix=curvature_reg_matrix, @@ -645,12 +662,12 @@ def patched( # which is correct, and is why the `d_np` rows stay comparable with the # library rows they are measured against. factor.clear() - stats: dict = {} - counts["last_stats"] = stats - return solver( + solve_stats: dict = {} + counts["last_stats"] = solve_stats + return injected( np.asarray(curvature_reg_matrix, dtype=float), np.asarray(data_vector, dtype=float), - stats=stats, + stats=solve_stats, ) patched.__wrapped__ = original diff --git a/scripts/misc/test/test_fixed_light_numpy_solvers.py b/scripts/misc/test/test_fixed_light_numpy_solvers.py index 28074bbed..188257225 100644 --- a/scripts/misc/test/test_fixed_light_numpy_solvers.py +++ b/scripts/misc/test/test_fixed_light_numpy_solvers.py @@ -190,6 +190,72 @@ def test_library_nnls_solve_reports_the_two_factorisations_it_spends(): assert stats["warm_start_fallback"] is False +@pytest.mark.parametrize("new_signature", [True, False]) +def test_numpy_injection_installs_against_the_old_and_new_library_signature( + monkeypatch, new_signature +): + """PyAutoArray #566 added ``solver``/``stats``; the injection must run on either side of it. + + The library passes ``solver=`` by name, so the wrapper must accept it on the numpy + path (and still dispatch to the injected kernel), and on the delegated JAX path + forward it only to a library that declares it. + """ + import types + + from autoarray.inversion.inversion import inversion_util + + calls = [] + if new_signature: + + def stub( + data_vector, + curvature_reg_matrix, + settings=None, + xp=np, + fingerprint=None, + factor=None, + solver="pdip", + stats=None, + ): + calls.append({"solver": solver, "stats": stats}) + return data_vector + + else: + + def stub( + data_vector, curvature_reg_matrix, settings=None, xp=np, fingerprint=None, factor=None + ): + calls.append({}) + return data_vector + + monkeypatch.setattr(inversion_util, "reconstruction_positive_only_from", stub) + + kernel_calls = [] + + def kernel(ZTZ, ZTx, *, stats): + kernel_calls.append(1) + return ZTx + + fake_jax = types.SimpleNamespace(__name__="jax.numpy") + stats = {} + with flns.numpy_solver_injected(kernel, label="stub") as counts: + call = inversion_util.reconstruction_positive_only_from + call(data_vector=np.ones(3), curvature_reg_matrix=np.eye(3), xp=np, solver="certified") + call( + data_vector=np.ones(3), + curvature_reg_matrix=np.eye(3), + xp=fake_jax, + solver="certified", + stats=stats, + ) + assert (counts["numpy"], counts["jax"], len(kernel_calls)) == (1, 1, 1) + if new_signature: + assert calls == [{"solver": "certified", "stats": stats}] + else: + assert calls == [{}], "the old library must not be handed solver/stats" + assert inversion_util.reconstruction_positive_only_from is stub + + # --------------------------------------------------------------------------- # 2. The jax-free contract # --------------------------------------------------------------------------- From 2b013f9d8544bd2b109cbc25ab7b4d519dcbdefc Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 23 Sep 2026 16:15:06 +0200 Subject: [PATCH 3/3] test(profiling): skip the #566-dependent fallback pin on a pre-#566 library lint.yml checks out PyAutoArray main, so asserting the #566 signature made this PR's CI red by construction until PyAutoArray#567 merges; skip there, assert on the new library. Co-Authored-By: Claude Opus 5.5 (1M context) --- scripts/misc/test/test_fixed_light_library.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scripts/misc/test/test_fixed_light_library.py b/scripts/misc/test/test_fixed_light_library.py index de5216567..6c6413870 100644 --- a/scripts/misc/test/test_fixed_light_library.py +++ b/scripts/misc/test/test_fixed_light_library.py @@ -246,9 +246,10 @@ def test_injection_installs_against_the_old_and_new_library_signature(monkeypatc def test_fallback_pins_the_library_pdip_solver(): """Route (d)'s fallback is the library's PDIP, whatever solver the library config selects.""" original = _library_positive_only() - assert "solver" in inspect.signature(original).parameters, ( - "this test needs the PyAutoArray #566 library on the path (source activate.sh)" - ) + if "solver" not in inspect.signature(original).parameters: + pytest.skip( + "requires PyAutoArray#566: reconstruction_positive_only_from has no `solver` kwarg" + ) seen = [] def spy(