Sparse triangular solve without QR factorization

[hanbin973]

I have a sparse triangular system to solve. Is it possible to solve this system directly by back/forward substitution without QR factorization in jax?

[dfm]

I think that lax.linalg.triangular_solve might do the trick for you?

[hanbin973]

I checked the function already but it only takes dense Jax arrays. My sparse system is too large to fill in all the zeros.

[Vlad382]

Hi!
I have been solving triangular system in JAX myself. There is a function I have come up with:

from jax import jit
from jax.lax import scan
import jax.numpy as jnp

@jit
def jspsolve_lower_triangular(A, b):
    x = jnp.zeros_like(b)    
    
    diag_edge_indx = jnp.argwhere(jnp.diff(A.indices) == 0, size=b.shape[0], fill_value=jnp.nan)[:, 0].astype(jnp.int32)
    xs_ = jnp.hstack([
        jnp.arange(x.shape[0])[:, None],
        diag_edge_indx[:, None]
    ])
    
    def f_(x_local, k):
        i, diag_ind = k
        x_i = x_local.at[jnp.where(A.indices[:, 0] == i, A.indices[:, 1], A.indices.shape[0])].get(mode='fill', fill_value=0)
        A_i = jnp.where(A.indices[:, 0] == i, A.data, 0)
        c = (b[i] - jnp.sum(A_i * x_i)) / (A.data[diag_ind] + 1e-16)
        x_local = x_local.at[i].set(c, mode='drop')
        return x_local, None
    
    x, _ = scan(f_, x, xs_)
    return x

It is compatible with JAX's transformations (grad, vmap, jit) and works with instances of the BCOO class from an experimental sparse module in JAX. However, this function operates on the entire vector of length NNZ during each substitution step. This behaviour is due to 1) JAX requires static shapes for traced arrays and 2) one cannot use traced arrays as indexes or to create arrays with. This was sufficient for my test on small problems, but I did not manage to make a reasonable function.