wip: N-D LinearOperator and cg

[lucascolley]

Building on scipy#23836

https://github.com/kokkos/kokkos-kernels/blob/develop/batched/sparse/impl/KokkosBatched_CG_Team_Impl.hpp

https://ieeexplore.ieee.org/document/10054414 section VI

[izaid]

If the new matvec can handle batched dims, do we need hstack or can it just be done all in one go?

[lucascolley]

@izaid to maintain compatibility with user-defined matvec functions which do not support batching, I think we do need to keep this hstack.

I think it could be done all in one go for batched matvec functions. But I think that would have to be opt-in in the API.

[izaid]

Yes! This looks good in general, basically what I was thinking. Will this error in all previous cases if someone passes an ND LinearOperator to an existing method. For instance, gmres should throw an error if the LinearOperator it receives has ndim != 2.

[lucascolley]

Will this error in all previous cases if someone passes an ND LinearOperator to an existing method.

that is also TODO 👍

[lucascolley]

Prototype working with array_api_strict and jax.numpy! (no JIT yet)

In [1]: import array_api_strict as xp; from scipy.sparse.linalg import cg, LinearOperator; import numpy as np

In [2]: def solve(N, batch, report_index=0, batched=False):
   ...:     rng = np.random.default_rng(0)
   ...:     M = rng.standard_normal((N, N))
   ...:     M = xp.asarray(M)
   ...:     reg = 1e-3
   ...:
   ...:     if batched:
   ...:         M = xp.broadcast_to(M[xp.newaxis, ...], (batch, *M.shape))
   ...:
   ...:     def matvec(x):
   ...:         return xp.squeeze(M.mT @ (M @ x[..., xp.newaxis]), axis=-1) + reg * x
   ...:
   ...:     shape = (batch, N, N) if batched else (N, N)
   ...:     A = LinearOperator(shape, matvec=matvec, dtype=xp.float64, xp=xp)
   ...:
   ...:     b = rng.standard_normal(N)
   ...:     b = xp.asarray(b)
   ...:
   ...:     if batched:
   ...:         b = xp.reshape(xp.arange(batch, dtype=xp.float64), (batch, 1)) * b
   ...:         x, info = cg(A, b, atol=1e-8, maxiter=5000)
   ...:         assert info == 0
   ...:         print(f"{x[report_index, ...]}")
   ...:     else:
   ...:         for i in xp.arange(batch, dtype=xp.float64):
   ...:             x, info = cg(A, i*b, atol=1e-8, maxiter=5000)
   ...:             assert info == 0
   ...:             if i == report_index:
   ...:                 print(x)
   ...:

In [3]: solve(5, 10, report_index=7)
Array([ 10.91985197,  -5.53737923,
        -6.96397906, -35.6473016 ,
        13.48931722], dtype=array_api_strict.float64)

In [4]: solve(5, 10, report_index=7, batched=True)
Array([ 10.91985197,  -5.53737923,
        -6.96397906, -35.6473016 ,
        13.48931722], dtype=array_api_strict.float64)

In [5]: import jax
   ...: jax.config.update("jax_enable_x64", True)

In [6]: import jax.numpy as xp

In [7]: solve(5, 10, report_index=7, batched=True)
[ 10.91985197  -5.53737923  -6.96397906 -35.6473016   13.48931722]

For JIT I think we will need to do some fancy registration of the linear operator classes as PyTrees, along the lines of https://docs.jax.dev/en/latest/_autosummary/jax.tree_util.register_pytree_node.html#jax.tree_util.register_pytree_node.

[izaid]

Oh that's great! Yes, might need to register LinearOperator as PyTrees, but not such a big deal I think?

[lucascolley]

I think the JIT will probably require a separate backend that uses JAX-specific things. In particular, for:

converged = xp_vector_norm(r, axis=-1) < atol
if xp.all(converged):
    return x, 0

I think we will need to use something from https://docs.jax.dev/en/latest/control-flow.html#structured-control-flow-primitives.

[lucascolley]

may need axis=-1 for the complex case too

[lucascolley]

should we support different atols for different systems? If not, should this be xp.min or xp.max?