Reuse repeated code between LOO (and likely MC) estimators

[pnkraemer]

Currently, the LOO estimators include cases like if num_samples >= num_rows: materialize_and_compute_exactly(). Since this duplicates a number of exceptions between the integrands, it would be useful to pull this code out of there, and, for example, move it to the estimator_leave_one_out implementation.

But this requires fiddling with the data structures in the stochtrace module, so it's relatively low priority for now.

[sethaxen]

For the LOO integrands in particular, much of the boilerplate could be eliminated while still supporting users writing their own integrands if we defined the following public types:

 class LOOContext(NamedTuple):
     """Everything an LOO estimator needs, computed once per `estimate(...)` call."""
     matvec: Callable
     params: tuple
     matvec_flat: Callable[[Array], Array]
     unflatten_input: Callable[[Array], Any]
     Omega: Array          # shape (n, num_samples)
     n: int
     num_samples: int


 class LOOBranches(NamedTuple):
     """The two ways to compute an LOO estimate, plus the cost that picks between them."""
     num_matvecs_per_sample: int
     exact: Callable[[Array], Array]    # exact(B_mat) -> per-sample result
     estimate: Callable[[], Array]      # estimate() -> per-sample result


 # A "LOO estimator" is any callable with this shape. All integrands return one of these callables.
 LOOEstimator = Callable[[LOOContext], LOOBranches]

It's the job of the estimator_leave_one_out[_...] methods to produce a method that generates the LOOContext and passes to the LOOEstimator to get the LOOBranches; based on the num_matvecs_per_sample compared to n, it decides which branch to take and returns the result pytree.

Then e.g. our implementations look like

 def leave_one_out_xtrace(*, apply_resphering: bool = True) -> LOOEstimator:
     def loo_estimator(ctx: LOOContext) -> LOOBranches:
         matvec_flat, Omega, num_samples = ctx.matvec_flat, ctx.Omega, ctx.num_samples

         def exact(B_mat):
             return np.ones((num_samples,), dtype=B_mat.dtype) * linalg.trace(B_mat)

         def estimate():
             matmat = func.vmap(matvec_flat, in_axes=-1, out_axes=-1)
             Y = matmat(Omega)
             Q, R = linalg.qr_reduced(Y)
             Z = matmat(Q)

             def _trace_exact():
                 ...

             def _trace_estimate():
                 ...

             return _cond_rank_deficient(R, num_samples, _trace_exact, _trace_estimate)

         return LOOBranches(num_matvecs_per_sample=2, exact=exact, estimate=estimate)

     return loo_estimator

while a user's combined integrand might look like

 def leave_one_out_combined(ctx: stochtrace.LOOContext) -> stochtrace.LOOBranches:
     xtrace = stochtrace.leave_one_out_xtrace()(ctx)
     xrownorms_squared = stochtrace.leave_one_out_xrownorms_squared()(ctx)
     k = max(xtrace.num_matvecs_per_sample, xrownorms_squared.num_matvecs_per_sample)

     def exact(B_mat):
         return {"trace": xtrace.exact(B_mat), "rownorms_squared": xrownorms_squared.exact(B_mat)}

     def estimate():
         return {"trace": trace_b.estimate(), "rownorms_squared": xrownorms_squared.estimate()}

     return stochtrace.LOOBranches(num_matvecs_per_sample=k, exact=exact, estimate=estimate)

This avoids the worst-case scenario where one integrand with many matvecs-per-sample materializes the operator to take the exact path while another computes an estimate, roughly doubling the work and making it impossible for CSE optimization.

I don't love this design, and it would be breaking and to me feels inconsistent with matfree's vibe, but what are your thoughts?