Issue with vmap when using lx.linear_solve on SparseMatrixOperator with multi-column RHS

[Dong555]

I'm encountering a ValueError when using the vmap functionality to map the lx.linear_solve operation across multiple columns of a right-hand side matrix (RHS) with a SparseMatrixOperator.

I expect the _solve_beta function when passed a SparseMatrixOperator and a multi-column RHS, to solve the linear system for each column of the RHS smoothly similarly to how it works when a MatrixLinearOperator is used. While the linear solve operation works fine for individual columns of y with a SparseMatrixOperator, it throws an error when attempting to solve for multiple columns using vmap, and produces the following error:
ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification ... (trimmed for brevity)

Here is a minimal working example provided specific to this problem:

from __future__ import annotations

import numpy as np
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import lineax as lx
from jax.experimental import sparse
from jaxtyping import Array, ArrayLike, Float

class SparseMatrixOperator(lx.AbstractLinearOperator):
    matrix: sparse.JAXSparse

    def __init__(self, matrix):
        self.matrix = matrix

    def mv(self, vector: ArrayLike):
        return sparse.sparsify(jnp.matmul)(
            self.matrix, vector, precision=lax.Precision.HIGHEST
        )

    def as_matrix(self) -> Float[Array, "n p"]:
        return self.matrix.todense()

    def transpose(self) -> "SparseMatrixOperator":
        return SparseMatrixOperator(self.matrix.T)

    def in_structure(self) -> jax.ShapeDtypeStruct:
        _, in_size = self.matrix.shape
        return jax.ShapeDtypeStruct((in_size,), self.matrix.dtype)

    def out_structure(self) -> jax.ShapeDtypeStruct:
        out_size, _ = self.matrix.shape
        return jax.ShapeDtypeStruct((out_size,), self.matrix.dtype)

@lx.is_negative_semidefinite.register(SparseMatrixOperator)
def _(op):
    return False

@lx.linearise.register(SparseMatrixOperator)
def _(op):
    return op

_multi_linear_solve = eqx.filter_vmap(lx.linear_solve, in_axes=(None, 1, None))

def _solve_beta(A: lx.AbstractLinearOperator, rhs: ArrayLike) -> Array:
    solver = lx.NormalCG(rtol=1e-6, atol=1e-6)
    out = _multi_linear_solve(A, rhs, solver)
    updated_beta = out.value.T

    return updated_beta

N, P, K = 10, 3, 2
X = jnp.asarray(np.random.normal(size=(N, P)))
y = jnp.asarray(np.random.normal(size=(N, K)))
solver = lx.NormalCG(rtol=1e-6, atol=1e-6)

# Not actually dense, but just need an instance...
G_sp = SparseMatrixOperator(sparse.BCOO.fromdense(X))
G_op = lx.MatrixLinearOperator(X)

# Works in the following cases
print(lx.linear_solve(G_sp, y[:, 0], solver).value)
print(lx.linear_solve(G_sp, y[:, 1], solver).value)
print(_solve_beta(G_op, y))

# Does not work in this case
print(_solve_beta(G_sp, y))

[patrick-kidger]

Thanks for the report! This looks to be a bug in JAX itself: jax-ml/jax#18244.

Sparse matrices still only have experimental support, so it looks like all the kinks haven't been worked out yet.

Since you're not batching over the operator, it looks like you can work around this by capturing the operator via closure, rather than passing it as an unbatched argument:

def _solve_beta(A: lx.AbstractLinearOperator, rhs: ArrayLike) -> Array:
    solver = lx.NormalCG(rtol=1e-6, atol=1e-6)
    out = jax.vmap(lambda b: lx.linear_solve(A, b, solver), in_axes=1)(rhs)
    updated_beta = out.value.T

    return updated_beta

Note that you do have to use jax.vmap rather than eqx.filter_vmap here -- the latter still ends up hitting a variant of the above bug in jax.experimental.sparse.BCOO.