Jax custom vjp for sparse array misinterprets size of the array

[aadityacs]

I am trying to define custom vjps in jax that works with jax sparse BCOO as inputs. However, I am running into an issue where the shapes that are interpreted for the backward for the sparse array is the number of specified elements rather than the shape of the matrix.

I am not sure how best to resolve this. Will be grateful for any pointers and help in this!

Below is a simple code that replicates the core issue:

#-------------imports----------------#
from typing import Tuple
import numpy as np
import jax
import jax.numpy as jnp
from jax.typing import ArrayLike 
from jax.experimental import sparse
#-------------custom vjp function----------------#
@jax.custom_vjp
def fn(matrix: sparse.BCOO | ArrayLike, eps: float
       ) -> Tuple[jnp.ndarray, jnp.ndarray]:
  """Some numpy based computation that needs to be made differentiable."""
  del eps # some other variables not computed gradient of.
  if not isinstance(matrix, ArrayLike):
    matrix = np.asarray(matrix.todense())
  
  evals = matrix.sum() * jnp.ones(matrix.shape[0]) # vector of (N,)
  evecs = matrix.sum() * jnp.ones_like(matrix) # Array of (N,N) 
  return evals, evecs


def fwd(
  matrix: sparse.BCOO,
  eps: float,
) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray, float]]:
  """Some numpy based computation that needs to be made differentiable."""
  evals, evecs = fn(matrix, eps)
  return (evals, evecs), (evals, evecs, eps)


def bwd(
  res: Tuple[jnp.ndarray, jnp.ndarray, float],
  grads: Tuple[jnp.ndarray, jnp.ndarray],
  ) -> Tuple[jnp.ndarray, None]:
  """Implements the backward calculation for `fwd`."""
  evals, evecs, eps = res
  grad_evals, grad_evecs = grads

  # grad_mtrx is some fn of evals, evecs, grad_evals, grad_evecs
  grad_mtrx = grad_evals[:, jnp.newaxis] + grad_evecs + jnp.sum(evals) + jnp.sum(evecs)
  return grad_mtrx, None


fn.defvjp(fwd, bwd)

#-------------some loss function----------------#
def loss_fn(matrix: sparse.BCOO | ArrayLike, eps: float):
  """Some scalar loss function that needs to be made differentiable."""
  evals, evecs = fn(matrix, eps) # call the custom autograd function
  return jnp.sum(evals) + jnp.sum(evecs)


#---------------test---------------------------#

# Define the data and indices for the sparse matrix
data = jnp.array([1., 2., 3., 4., 5.,])
indices = jnp.array([[0, 0], [0, 1], [1, 0], [1, 1], [2,2]])

sparse_matrix = sparse.BCOO((data, indices), shape=(3, 3))
dense_matrix = sparse_matrix.todense()

eps = 1.e-3
# This works when passing the dense matrix
v, grad_v = jax.value_and_grad(loss_fn, allow_int=True)(dense_matrix, eps)

# This fails when passing the sparse matrix.
# The function interprets the number of specified elements (5) as the size
# of the sparse matrix instead of the actual size (3X3)

v, grad_v = jax.value_and_grad(loss_fn, allow_int=True)(sparse_matrix, eps)

[jakevdp]

It looks like the issue is your bwd rule is returning a dense matrix for sparse input, when it should be returning a sparse matrix that matches the sparsity pattern of the input.

[aadityacs]

Thanks for the reply @jakevdp

In my case, I don't think my operations preserve the sparsity. I am wondering how to go about in this scenario.

To provide more context, I am writing jax wrappers for scipy eigsh. The returned eigenvectors are usually dense and my understanding is that the output of bwd will also be dense.