[sparse] Wrong jacobian computation for bcoo_spdot_general

It seems that the Jacobian computed by bcoo_spdot_general, invoked through the operator @ does not organize the values in the correct order.

In order to show the bug, I have written the following code.

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

A_indices = jnp.array([[0, 1], [0, 2], [1, 1], [1, 2], [1, 0]])
A_values = jnp.array([-2.0, 1.0, -1.0, 0.5, 2.0])
A_shape = (2, 3)

B_indices = jnp.array([[0, 2], [2, 1], [0, 3], [1, 3], [1, 0], [0, 0]])
B_values = jnp.array([10.0, 100.0, 1000.0, -5.0, -50.0, -500.0])
B_shape = (3, 4)

def sp_sp_product(v1, v2):
    sparse1 = BCOO((v1, A_indices), shape=A_shape)
    sparse2 = BCOO((v2, B_indices), shape=B_shape)
    return (sparse1 @ sparse2).todense()

def sp_dense_product(v1, v2):
    sparse1 = BCOO((v1, A_indices), shape=A_shape)
    dense2 = BCOO((v2, B_indices), shape=B_shape).todense()
    return sparse1 @ dense2

sp_sp_product_jac = jax.jacfwd(sp_sp_product, argnums=1)(A_values, B_values)
sparse_dense_product_jac = jax.jacfwd(sp_dense_product, argnums=1)(A_values, B_values)

assert jnp.array_equal(
    sp_sp_product_jac, sparse_dense_product_jac
), "The two Jacobian should be the same"

sparse_dense_product_jac is correct (as it can be verified analitically), while sp_sp_product_jac is incorrect since it should be identical to sparse_dense_product_jac but it is not. The shape of sp_sp_product_jac in the expected one and the values in sp_sp_product_jac seems correct but they are not placed in the expected location. Probably there is a problem in the indices of the sparse jacobian matrix.

[jakevdp]

I don’t think this will require any deep fix, we just need to not use a problematic version of sum_duplicates in the bcoo_spdot_general implementation. I think the solution will be to split the current sum_duplicates utility into a couple utilities, or perhaps one with several options: sum_duplicates, sort_indices, and eliminate_zeros. The first two are fine for autodiff, the third should error on autodiff. Then in functions that we want to be differentiable, we don’t call the third.

[jakevdp]

Put another way, this issue has brought up a more general concern that I hadn’t considered before: sparse autodiff is fundamentally incompatible with any operation in which the output indices depend on the values in the data, because that may break the property that primals and tangents must share indices.