fix(profiling): accept PyAutoArray's new solver/stats kwargs in the solver injections (PyAutoArray#566)

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

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(

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,96 @@ 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()
+    if "solver" not in inspect.signature(original).parameters:
+        pytest.skip(
+            "requires PyAutoArray#566: reconstruction_positive_only_from has no `solver` kwarg"
+        )
+    seen = []
+
+    def spy(
+        data_vector,
+        curvature_reg_matrix,
+        settings=None,
+        xp=np,
+        fingerprint=None,
+        factor=None,
+        solver="<absent>",
+        stats="<absent>",
+    ):
+        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
 # ---------------------------------------------------------------------------

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
 # ---------------------------------------------------------------------------