Specify custom_jvps for sparse matrices

[SNMS95]

Hi guys,

I want to use an external sparse linear solver alongside JAX, which requires a matrix-vector product function and an RHS vector.
I know that I have to use pure_callback + custom_jvp on this solver to enable higher order autodiff. This can be done. However, the issue is the sparse-matrix creation.

For saving memory, I have to create the sparse matrix from ordinary numpy arrays.

def create_sparse_matrix(V, I, J, shape) # ----> They are all numpy arrays [Consider this as host call]
    A_sp_scipy = scipy.sparse.csr_array((V, (I, J)),
                                         shape=shape)
    # Need to cache the scipy version for a later use
    A_sp = BCOO.from_scipy_sparse(A_sp_scipy).sort_indices()
    return A_sp

@jax.custom_jvp
def creation(V, I, J, size):
    # Convert to JAX arrays needed ? Don't think so
    # How to specify the shape for sparse matrices??????????????????
    result_shape = jax.ShapeDtypeStruct(shape=(size, size), dtype=np.float64)
    Asp = jax.pure_callback(create_sparse_BC00, result_shape, V, I, J, size)
    return Asp

@creation.def_jvp
_creation_jvp(primals, tangents):
    V, I, J, size = primals
    V_dot, _, _, _ = tangents # tangents do not exist for integers
    out = creation(V, I, J, size)
    out_dot = creation(V_dot, I, J, size)
    return out, out_dot

The main question is "How do I specify the result_shape for sparse matrices in the pure_callback?"
It would be nice, if you point out if there is a something wrong with this approach!

[nopror]

ef create_sparse_matrix(V, I, J, shape) # ----> They are all numpy arrays [Consider this as host call]
A_sp_scipy = scipy.sparse.csr_array((V, (I, J)),
shape=shape)
# Need to cache the scipy version for a later use
A_sp = BCOO.from_scipy_sparse(A_sp_scipy).sort_indices()
return A_sp

@jax.custom_jvp
def creation(V, I, J, size):
# Convert to JAX arrays needed ? Don't think so
# How to specify the shape for sparse matrices??????????????????
result_shape = jax.ShapeDtypeStruct(shape=(size, size), dtype=np.float64)
Asp = jax.pure_callback(create_sparse_BC00, result_shape, V, I, J, size)
return Asp

@creation.def_jvp
_creation_jvp(primals, tangents):
V, I, J, size = primals
V_dot, _, _, _ = tangents # tangents do not exist for integers
out = creation(V, I, J, size)
out_dot = creation(V_dot, I, J, size)
return out, out_dot