Sparse matrix operations and jaxlib version

[pvasired]

I am hoping to do sparse matrix operations on some data because I am finding that dense matrix operations are taking too long. In particular, I want to use the jax.experimental.sparse.linalg.spsolve() because I am inverting a large, sparse matrix using jax.numpy.linalg.solve() and finding that it takes prohibitively long. I tried following the nice documentation here for using sparse matrices in JAX: [https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html#module-jax.experimental.sparse]

But I found that only the CSR data structure can be used with spsolve(). For one, it seems like the BCOO data structure is better supported -- will BCOO become useable in spsolve() any time soon?

Furthermore, I'm having trouble even getting spsolve() to run using GPU. I had to install a specific (and old) version of jaxlib a while back in order to be compatible with the CUDA drivers on my computing server. However, I get an error message saying ValueError: spsolve requires jaxlib version 86 or above. But I can't update my jaxlib version or I'll lose compatibility with my GPUs.

I made a new conda environment to test this out, using the most recent versions of jaxlib and jax, and sure enough I can now call spsolve(), but I get the warning that No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Any ideas on what I can try here? It seems like a very tricky spot to be in where my CUDA drivers are too old to be used with spsolve(), and my current implementation without sparse operations is looking to be too slow to be used in practice.

[pvasired]

I also just saw in a recent post that while sparse matrices may improve memory costs, they might not actually improve runtime significantly. In which case, are there any optimizations I can perform on my code to get it to run faster? The objective function I'm using in a training loop (using the optax library) is given below. These matrices can get pretty large, and I think the costliest line is the jnp.linalg.solve(). I had to play some games with avoiding creating large diagonal matrices and avoiding explicit trace calculation since the matrices are so big that they can't even be stored in my GPU memory.

@jax.jit
def fisher_loss_max(probs_vec, transform_mat, jac_full, trials):
    """
    Compute the Fisher loss across the entire array.

    Parameters:
    probs_vec (jnp.DeviceArray): The flattened array of probabilities across all meaningful
                            (cell, pattern) combinations
    transform_mat (jnp.DeviceArray): The transformation matrix to convert the trials array
                                to the transformed trials array for multiple cells on
                                the same pattern
    jac_full (jnp.DeviceArray): The full precomputed Jacobian matrix
    trials (jnp.DeviceArray): The input trials vector to be optimized

    Returns:
    loss (float): The whole array Fisher information loss
    """
    p_model = jnp.clip(probs_vec, a_min=1e-5, a_max=1-1e-5) # need to clip these to prevent
                                                            # overflow errors
    t = jnp.dot(transform_mat, trials).flatten()
    I_p = t / (p_model * (1 - p_model))

    # Avoiding creating the large diagonal matrix and storing in memory
    I_w = jnp.dot((jac_full.T * I_p), jac_full) / len(p_model)
    
    # Avoiding multiplying the matrices out and calculating the trace explicitly
    sum_probs = jnp.sum(jnp.multiply(jac_full.T, jnp.linalg.solve(I_w, jac_full.T)), axis=0)
    sum_cells = jnp.reshape(sum_probs, (-1, trials.shape[1])).sum(axis=-1)

    return jax.scipy.special.logsumexp(sum_cells)

I'm not sure if people here have much experience using optax, but I thought I'd attach the training loop as well in case anyone has thoughts on how to speed it up. I've trivially found that using other optimizers than AdamW can speed things up, but AdamW seems to be the most stable and reliable so I've stuck with it.

def optimize_fisher_array(jac_full, probs_vec, transform_mat, T_prev, T, reg=None, 
                          step_size=1, n_steps=100, T_budget=5000, verbose=True):
    """
    Fisher optimization loop using optax and AdamW optimizer.

    Parameters:
    jac_full (jnp.DeviceArray): Full precomputed Jacobian matrix
    probs_vec (jnp.DeviceArray): Flattened array of all probabilities
                                 for all non-degenerate (cell, pattern)
                                 combinations
    transform_mat (jnp.DeviceArray): transformation matrix to convert trials
                                     array into correct shape matrix for 
                                     multiple cells on the same pattern
    T_prev (jnp.DeviceArray): The previously sampled trials array
    T (jnp.DeviceArray): The initialization for the to-be-optimized trials

    Returns:
    losses (np.ndarray): An array of losses per iteration of the optimization routine
    T (jnp.DeviceArray): The optimized trials matrix 
    """
    # Initialize the optimizer
    optimizer = optax.adamw(step_size)
    opt_state = optimizer.init(T)

    if reg is None:
        init_function = fisher_loss_max(probs_vec, transform_mat, jac_full, T_prev + jnp.absolute(T))
        reg = init_function / 20000 # 20000 worked, 100000 too large

    # Update function for computing the gradient
    @jax.jit
    def update(jac_full, probs_vec, transform_mat, T_prev, T):
        # Adding special l1-regularization term that controls the total trial budget
        fisher_lambda = lambda T, jac_full, probs_vec, transform_mat, T_prev: fisher_loss_max(probs_vec, transform_mat, jac_full, T_prev + jnp.absolute(T)) + reg * jnp.absolute(jnp.sum(jnp.absolute(T)) - T_budget)
        grads = jax.grad(fisher_lambda)(T, jac_full, probs_vec, transform_mat, T_prev)
        
        return grads
    
    losses = []
    for step in range(n_steps):
        if verbose:
            print(step)

        # Update the optimizer
        grads = update(jac_full, probs_vec, transform_mat, T_prev, T)
        updates, opt_state = optimizer.update(grads, opt_state, params=T)
        T = optax.apply_updates(T, updates)
        
        # If desired, compute the losses and store them
        if verbose:
            loss = fisher_loss_max(probs_vec, transform_mat, jac_full, T_prev + jnp.absolute(T))
            loss_tuple = (loss, jnp.sum(jnp.absolute(T)), loss + reg * jnp.absolute(jnp.sum(jnp.absolute(T)) - T_budget),
                            reg * jnp.absolute(jnp.sum(jnp.absolute(T)) - T_budget))
            print(loss_tuple)
            losses += [loss_tuple]

    return np.array(losses), T

[jakevdp]

But I found that only the CSR data structure can be used with spsolve().

That's correct. This function lowers to a GPU primitive that only supports CSR input.

For one, it seems like the BCOO data structure is better supported

That's correct. Most of our existing JAX sparse functionality is built for BCOO. But BCSR is getting more complete, and that should work with spsolve(), so that's worth a try.

will BCOO become useable in spsolve() any time soon?

Unlikely, because the GPU lowering for spsolve requires CSR.

Furthermore, I'm having trouble even getting spsolve() to run using GPU.

Yeah, CUDA versioning is a real pain. Cusparse/cusolver in particular seems to be pretty unstable between releases, which makes it difficult for a library like JAX to target effectively because things break with little warning, and standing up a comprehensive CI test matrix across CUDA/cusparse/cusolver/cudnn/hardware versions is too expensive to justify.

Any ideas on what I can try here? It seems like a very tricky spot to be in where my CUDA drivers are too old to be used with spsolve(), and my current implementation without sparse operations is looking to be too slow to be used in practice.

No, no real suggestions. This is why JAX sparse is still in jax.experimental – if what's there is useful, then go for it! But if not, we don't have cycles to provide full support, particularly given the CUDA challenges mentioned above.

Hope that helps!

[pvasired]

Got it, thank you for the help! Not a whole lot to be done here it seems, but I appreciate the input.