Is there a way to jit sparse linear solvers (jax.scipy.sparse.linalg) with A expressed as an operator?

[peterdsharpe]

Hey all, I'm interested in solving sparse linear systems in JAX.

The docs for the JAX sparse linear solver API (shared across all solvers in jax.scipy.sparse.linalg) indicate that the A parameter should be either:

• the A-matrix itself, or
• a callable linear operator for the A-matrix, given in the form lambda x: A @ x

In my application, providing a function that computes a matrix-vector product A @ x is much cheaper than directly computing the matrix A itself. However, if I provide A as a function, I'm no longer able to jit the sparse linear solve, which makes the solve much slower.

Here's a minimal failing example:

import jax
import jax.numpy as np
from jax.scipy import sparse
from jax import random

key = random.PRNGKey(seed=0)

N = 3

A = random.uniform(key=key, shape=(N, N))
b = random.uniform(key=key, shape=(N,))

A_operator = lambda v: A @ v # Just an example; for the real application we compute this more efficiently

##### This works!
x, _ = sparse.linalg.gmres(
    A=A_operator,
    b=b
)

##### But, add a `jit` and it no longer works!
try:
    x, _ = jax.jit(sparse.linalg.gmres)(
        A=A_operator,
        b=b
    )
except TypeError as e:
    print("Didn't work!")
    print(e)

The resulting error has to do with the input type to sparse.linalg.gmres being a function.

The relevant part of the error traceback is here:

Traceback (most recent call last):
  File "C:/GitHub/jax-opti/scratches/sparse_linsolve_jit_bug.py", line 18, in <module>
    x_sparse, _ = jax.jit(sparse.linalg.gmres)(
  File "C:\Users\peter\miniconda3\lib\site-packages\jax\_src\traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "C:\Users\peter\miniconda3\lib\site-packages\jax\_src\api.py", line 425, in cache_miss
    _check_arg(arg)
  File "C:\Users\peter\miniconda3\lib\site-packages\jax\_src\api.py", line 2928, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
TypeError: Argument '<function <lambda> at 0x000001D7655A8C10>' of type <class 'function'> is not a valid JAX type.

So, the question: Is there a way to jit sparse linear solvers (jax.scipy.sparse.linalg) with A expressed as an operator?

---

Miscellaneous details:

Verified that this occurs on both Windows 11 (manually built, CUDA backend) and Ubuntu 20.04 (on WSL, with CUDA backend); both tests run with Python 3.9.7, CUDA 11.6, Jax 0.3.1 and Jaxlib 0.3.0.

I can confirm that in addition to this not working with the gmres solver, it also doesn't work with cg, bistabcg, ...

[nicholasjng]

I think that's what static_argnums is for in the jax.jit API: marking certain arguments (my rule of thumb: not JAX-array-coercible ones like strings, callables etc.) constant ("static") for a JIT-compiled function call. These will then need to stay the same for as long as possible to avoid cache misses, which trigger full recompiles of the function.

So, in your example, if A is the n-th argument to the jax.scipy.sparse.linalg.gmres function, calling jax.jit(gmres, static_argnums=(n,)) should fix this.

[peterdsharpe]

Gotcha, thanks! This worked, and makes sense in retrospect.

Now that providing the linear operator is working, I'm running into another strange issue: the linear solver algorithms (e.g., jax.scipy.sparse.linalg.gmres itself) are significantly slower (~factor of 2) when given linear operators compared to when I give the A-matrix directly.

This is unexpected, because it looks like if you given gmres an array, it converts it into an operator internally anyway. This slowdow-when-given-an-operator happens both with and without JIT everywhere, and .block_until_ready() is called at the end of the timing context.

Anyway, I'll keep investigating - thanks for the help!

[YouJiacheng]

@peterdsharpe I cannot reproduce your problem(I made a mistake before)

def timer(f: Callable[[], Any]):
    from time import time
    f() # warmup

    t = time()
    for _ in range(3):
        f()
    print((time() - t) / 3)

from jax.scipy import sparse
from jax import random
from jax.tree_util import Partial
key = random.PRNGKey(seed=0)
N = 30
A = random.uniform(key=key, shape=(N, N))
b = random.uniform(key=key, shape=(N,))

A_operator = lambda v: A @ v
f = jax.jit(sparse.linalg.gmres)
g = jax.jit(sparse.linalg.gmres, static_argnums=0)
wraped = Partial(A_operator)
timer(lambda: jax.block_until_ready(f(A=wraped, b=b)))
timer(lambda: jax.block_until_ready(f(A=A, b=b)))
timer(lambda: jax.block_until_ready(g(A=wraped, b=b)))
timer(lambda: jax.block_until_ready(g(A=A_operator, b=b)))
# cannot pass A to g since array is not hashable 

0.6391735871632894
0.6313837369283041
0.6398694515228271
0.6331314245859782

Please check if you perform a warmup call before the speed measurement, since jit compiling is slow.
(but you said the without JIT it still slower, hmmm)

def timer_one(f: Callable[[], Any]):
    from time import time
    t = time()
    f()
    print(time() - t)
timer_one(lambda: f.lower(A=wraped, b=b).compile())
timer_one(lambda: f.lower(A=A, b=b).compile())
timer_one(lambda: g.lower(A=wraped, b=b).compile())
timer_one(lambda: g.lower(A=A_operator, b=b).compile())

0.7742102146148682
0.7620875835418701
0.7628335952758789
0.7621660232543945

[YouJiacheng]

Besides using static_argnums, you can also use jax.tree_util.Partial and pass zero arguments.

from jax.tree_util import Partial
x, _ = jax.jit(sparse.linalg.gmres)(
        A=Partial(A_operator),
        b=b
)

note jax.tree_util.Partial is different from functools.partial
BTW, I think you can always include sparse.linalg.gmres and the operator(closure with params) in a big function, and jit this big function.