Vmap batching of sparse linear system equation solver

[RodrigoCoelho7]

I was wondering if there is anyway I can batch (using vmap) a direct solver of a linear system of equations Ax=B where A is sparse?

Unfortunately, the jax.experimental.sparse.linalg.psolve function is not compatible with vmap.