Implementing a custom JVP for a sparse/iterative solver

[JanLuca]

I want to implement the JVP for a sparse eigenvalue solver as for example submitted in PR #32880 .

There I encounter the problem that for sparse/iterative solvers normally not a dense array is supplied but instead a function f(x) implementing a matrix-vector. I struggle now how implement a custom JVP for example in the case where function itself would have a well-defined derivative (for example a f(x) = A @ B @ x operation) which is not constructed due to performance/memory reasons.
How can I obtain a function df(x) implementing the action of the differential df? Can this be formulated as a general JVP function for the sparse solver or has one to do it at a higher level where the solver is called?

[jakevdp]

I don't think you can use custom_jvp to define the jvp of a function which is defined in this way: JAX's autodiff operations are defined with respect to functions which accept arrays as arguments. I'm not entirely sure how that could be made to work with a function that accepts a function as an argument.

Probably you'll have to define the autodiff behavior at a lower level, where the inputs and outputs are arrays rather than functions that operate on arrays.

[jakevdp]

Another possibility: if the user-specified function is a pytree, then I think custom_jvp could work at the level of the main API.