Computing gradients involving sparse arrays

[SNMS95]

Hi everyone.

I am trying to create a simple code as follows:

import jax
import jax.numpy as jnp
from jax.experimental import sparse

@jax.jit
def find_energy(x):
    # Sparse stiffness matrix (5x5)
    row_indices = jnp.array([0, 0, 1, 1, 2, 2, 3, 3, 4])
    col_indices = jnp.array([0, 1, 0, 1, 1, 2, 2, 3, 3])
    values = jnp.array([10.0, 2.0, 2.0, 8.0, -1.0, 5.0, 1.0, 4.0, 3.0])

    # Create the sparse matrix K in BCOO format
    K = x * sparse.BCOO((values, jnp.vstack([row_indices, col_indices]).T), shape=(nDof, nDof))
    P = jnp.eye(nDof)[free] 
    P_sp = sparse.BCOO.fromdense(P, nse=len(free), n_batch=0)
    K_free = (P_sp @ K) @ P_sp.T
    K_bcsr = sparse.BCSR.from_bcoo(K_free)
    fixed_rhs = F[free]
    sol_red = sparse.linalg.spsolve(K_bcsr.data, K_bcsr.indices, K_bcsr.indptr, fixed_rhs, tol=1e-06, reorder=1)
    full_sol = jnp.zeros(nDof)
    full_sol = full_sol.at[free].set(sol_red)
    c_value = jnp.dot(full_sol, F)
    return c_value

# Set up a small example
nDof = 5  # total number of degrees of freedom

# Load vector
F = jnp.array([1.0, 0.0, 0.0, 0.0, 0.0])

# Assume that the first and last DOFs are fixed
fixed = jnp.array([1, 4])
free = jnp.setdiff1d(jnp.arange(nDof), fixed)

# Design variable, usually Len(x) >>> 1
x = 1.0

# Call the function
energy = find_energy(x)
print("Compliance (Energy):", energy)
jax.grad(find_energy)(x) # Gradient of the energy with respect to x

This causes the error: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'bcoo_spdot_general' not implemented

If I try forward-mode (which is not preferred since x should have a large dimension but output is always scalar), I get an error in the solver(NotImplementedError: Batching rule for 'csr_matvec' not implemented). I want this to scale to larger problems, that is why I am using sparse matrices.

How can I fix this issue ?

1. Taking the gradient through the solver and multiplications
2. Is setting n_batch=0 correct? [I set it since K.n_batch=0]

[jakevdp]

You're finding out why all of this is under experimental 😁

Even if you got the gradients working, the sparse-sparse matmul code is terribly inefficient (see the discussion at #17251): I'd suggest finding a way to write your code without sparse-sparse matrix products, and that would solve both problems.

[SNMS95]

Thanks for the quick reply. If I just try to use P instead of its sparse version, after multiplication, we get an Array. However, this is not nice for memory right?

This is for a FEM code so K might be like 50k x 50k or something. What would be an alternative ?

[jakevdp]

Right: sparse @ dense = dense, because there's no guarantee that the output is sparse in general. You might be able to construct the output sparse matrix using something other than matrix multiplication.

[SNMS95]

Ahh. I was just using the multiplication to get around indexing [selecting a few rows and columns] since that is known to be an issue in XLA : #16291

This is quite difficult!
Would you have any suggestions ?

Also, thanks for all the hand-holding!

[jakevdp]

No, I don't have any suggestions. Unfortunately XLA is not a sparse-aware compiler, so any sparse computations lowered to XLA have to be done carefully.