From d2d89f17065b902683e776c052b7f7015fdfe4a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicholas=20Kr=C3=A4mer?= <nicokraemer@gmx.net>
Date: Mon, 18 Dec 2023 16:18:41 +0100
Subject: [PATCH 1/2] Change order of doctests/unittests in Makefile

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index b87ec40..50a4aaf 100644
--- a/Makefile
+++ b/Makefile
@@ -7,8 +7,8 @@ lint:
 	pre-commit run --all-files
 
 test:
-	python -m doctest README.md
 	pytest -x -v
+	python -m doctest README.md
 
 
 clean-preview:

From ba18b491eebac3945422b45fe11a6208153b74c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicholas=20Kr=C3=A4mer?= <nicokraemer@gmx.net>
Date: Mon, 18 Dec 2023 16:34:21 +0100
Subject: [PATCH 2/2] Implement optional matvec-parameters for Hutchinson's
 estimator

---
 matfree/hutchinson.py | 24 ++++++++++++------------
 matfree/slq.py        | 10 +++++-----
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/matfree/hutchinson.py b/matfree/hutchinson.py
index 1908bd2..aea93c2 100644
--- a/matfree/hutchinson.py
+++ b/matfree/hutchinson.py
@@ -19,8 +19,8 @@ def integrand_diagonal(matvec, /):
     where ``*args_like`` is an argument of the sampler.
     """
 
-    def integrand(v, /):
-        Qv = matvec(v)
+    def integrand(v, *parameters):
+        Qv = matvec(v, *parameters)
         v_flat, unflatten = tree_util.ravel_pytree(v)
         Qv_flat, _unflatten = tree_util.ravel_pytree(Qv)
         return unflatten(v_flat * Qv_flat)
@@ -31,8 +31,8 @@ def integrand(v, /):
 def integrand_trace(matvec, /):
     """Construct the integrand for estimating the trace."""
 
-    def integrand(v, /):
-        Qv = matvec(v)
+    def integrand(v, *parameters):
+        Qv = matvec(v, *parameters)
         v_flat, unflatten = tree_util.ravel_pytree(v)
         Qv_flat, _unflatten = tree_util.ravel_pytree(Qv)
         return linalg.vecdot(v_flat, Qv_flat)
@@ -43,8 +43,8 @@ def integrand(v, /):
 def integrand_trace_and_diagonal(matvec, /):
     """Construct the integrand for estimating the trace and diagonal jointly."""
 
-    def integrand(v, /):
-        Qv = matvec(v)
+    def integrand(v, *parameters):
+        Qv = matvec(v, *parameters)
         v_flat, unflatten = tree_util.ravel_pytree(v)
         Qv_flat, _unflatten = tree_util.ravel_pytree(Qv)
         trace_form = linalg.vecdot(v_flat, Qv_flat)
@@ -57,8 +57,8 @@ def integrand(v, /):
 def integrand_frobeniusnorm_squared(matvec, /):
     """Construct the integrand for estimating the squared Frobenius norm."""
 
-    def integrand(vec, /):
-        x = matvec(vec)
+    def integrand(vec, *parameters):
+        x = matvec(vec, *parameters)
         v_flat, unflatten = tree_util.ravel_pytree(x)
         return linalg.vecdot(v_flat, v_flat)
 
@@ -71,8 +71,8 @@ def integrand_trace_moments(matvec, moments, /):
     def moment_fun(x):
         return tree_util.tree_map(lambda m: x**m, moments)
 
-    def integrand(vec, /):
-        x = matvec(vec)
+    def integrand(vec, *parameters):
+        x = matvec(vec, *parameters)
         v_flat, unflatten = tree_util.ravel_pytree(vec)
         x_flat, _unflatten = tree_util.ravel_pytree(x)
         fx = linalg.vecdot(x_flat, v_flat)
@@ -139,9 +139,9 @@ def hutchinson(integrand_fun, /, sample_fun, stats_fun=np.mean):
 
     """
 
-    def sample(key):
+    def sample(key, *parameters):
         samples = sample_fun(key)
-        Qs = func.vmap(integrand_fun)(samples)
+        Qs = func.vmap(lambda vec: integrand_fun(vec, *parameters))(samples)
         return tree_util.tree_map(lambda s: stats_fun(s, axis=0), Qs)
 
     return sample
diff --git a/matfree/slq.py b/matfree/slq.py
index e9d6ba3..5303a9a 100644
--- a/matfree/slq.py
+++ b/matfree/slq.py
@@ -23,12 +23,12 @@ def integrand_slq_spd(matfun, order, matvec, /):
     This function assumes a symmetric, positive definite matrix.
     """
 
-    def quadform(v0, /):
+    def quadform(v0, *parameters):
         v0_flat, v_unflatten = tree_util.ravel_pytree(v0)
 
         def matvec_flat(v_flat):
             v = v_unflatten(v_flat)
-            Av = matvec(v)
+            Av = matvec(v, *parameters)
             flat, unflatten = tree_util.ravel_pytree(Av)
             return flat
 
@@ -83,12 +83,12 @@ def integrand_slq_product(matfun, depth, matvec, vecmat, /):
     Here, "product" refers to $X = A^\top A$.
     """
 
-    def quadform(v0, /):
+    def quadform(v0, *parameters):
         v0_flat, v_unflatten = tree_util.ravel_pytree(v0)
 
         def matvec_flat(v_flat):
             v = v_unflatten(v_flat)
-            Av = matvec(v)
+            Av = matvec(v, *parameters)
             flat, unflatten = tree_util.ravel_pytree(Av)
             return flat, tree_util.partial_pytree(unflatten)
 
@@ -97,7 +97,7 @@ def matvec_flat(v_flat):
 
         def vecmat_flat(w_flat):
             w = w_unflatten(w_flat)
-            wA = vecmat(w)
+            wA = vecmat(w, *parameters)
             return tree_util.ravel_pytree(wA)[0]
 
         # Decompose into orthogonal-bidiag-orthogonal