Repository: google/jax

Issue:
Differentiating BCOO sparse matrix constructor produces a dense matrix
### Description

Unless I am mistaken, differentiating the BCOO constructor appears to produce dense matrices. Here's an MWE:

```python
import jax 
import jax.numpy as jnp
from jax.experimental.sparse import BCOO

n = 1000
m = 777
p = 11
nonzeroes = jnp.stack((jnp.arange(0,m), jnp.arange(p,m+p)), axis=-1)

def quadratic(p):
    x1 = nonzeroes[...,0]
    x2 = nonzeroes[...,1]
    values = (x1 - x2)**2 / p
    matrix = BCOO((values, nonzeroes), shape=(n,n), unique_indices=True)
    return jnp.ones(n).T @ matrix @ jnp.ones(n)
    
graph = jax.xla_computation(jax.grad(quadratic))(1.0)

with open("t.dot", "w") as f:
    f.write(graph.as_hlo_dot_graph())
```

I've attached a computational graph: [t.pdf](https://github.com/google/jax/files/10804965/t.pdf). The problem is `dot.28` (dark blue) followed by `gather.57`, which appears to assemble a 1000x1000 matrix, then grab the parts of it that are zero.

### What jax/jaxlib version are you using?

0.4.4

### Which accelerator(s) are you using?

N/A

### Additional system info

N/A

### NVIDIA GPU info

N/A


Discussion:
I don't think any dense matrices are produced on the JAX side of things. For example, here is the jaxpr of your differentiated function, which does not show any `1000,1000` variable being produced:
```python
print(jax.make_jaxpr(jax.grad(quadratic))(1.0))
```
```
{ lambda a:i32[777,2]; b:f32[]. let
    c:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    d:i32[777] = gather[
      dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,))
      fill_value=None
      indices_are_sorted=True
      mode=GatherScatterMode.PROMISE_IN_BOUNDS
      slice_sizes=(777, 1)
      unique_indices=True
    ] a c
    e:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
    f:i32[777] = gather[
      dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,))
      fill_value=None
      indices_are_sorted=True
      mode=GatherScatterMode.PROMISE_IN_BOUNDS
      slice_sizes=(777, 1)
      unique_indices=True
    ] a e
    g:i32[777] = sub d f
    h:i32[777] = integer_pow[y=2] g
    i:f32[777] = convert_element_type[new_dtype=float32 weak_type=True] h
    j:f32[777] = div i b
    k:f32[] = integer_pow[y=-2] b
    l:f32[1000] = broadcast_in_dim[broadcast_dimensions=() shape=(1000,)] 1.0
    m:f32[777] = convert_element_type[new_dtype=float32 weak_type=False] j
    n:f32[1000] = bcoo_dot_general[
      dimension_numbers=(((0,), (0,)), ((), ()))
      lhs_spinfo=BCOOInfo(shape=(1000, 1000), indices_sorted=False, unique_indices=True)
    ] m a l
    o:f32[1000] = broadcast_in_dim[broadcast_dimensions=() shape=(1000,)] 1.0
    _:f32[] = dot_general[
      dimension_numbers=(((0,), (0,)), ((), ()))
      precision=None
      preferred_element_type=None
    ] n o
    p:f32[1000] = dot_general[
      dimension_numbers=(((), ()), ((), ()))
      precision=None
      preferred_element_type=None
    ] 1.0 o
    q:f32[777] = broadcast_in_dim[broadcast_dimensions=() shape=(777,)] 1.0
    _:f32[777] r:i32[777,2] = bcoo_transpose[
      permutation=(1, 0)
      spinfo=BCOOInfo(shape=(1, 1), indices_sorted=False, unique_indices=False)
    ] q a
    s:f32[777] = bcoo_dot_general_sampled[
      dimension_numbers=(((), ()), ((), ()))
    ] p l r
    t:f32[777] _:i32[777,2] = bcoo_transpose[
      permutation=(1, 0)
      spinfo=BCOOInfo(shape=(1, 1), indices_sorted=False, unique_indices=False)
    ] s r
    u:f32[777] = convert_element_type[new_dtype=float32 weak_type=True] t
    v:f32[777] = mul u k
    w:f32[777] = mul v i
    x:f32[] = reduce_sum[axes=(0,)] w
    y:f32[] = neg x
  in (y,) }
```
That said, what might be happening is that you're hitting this TODO: https://github.com/google/jax/blob/7e001d842e1cbfdbf991a4bd9b236012cc40fba4/jax/experimental/sparse/bcoo.py#L1091-L1092

`dot_general_sampled` is used in the transpose rule of `dot_general`, which requires extracting sparse indices from a dense dot product. There are probably more efficient ways to implement this, thus the TODO
One way around this would be to use forward-mode autodiff, which does not use the transpose rule.
Thanks for your quick response! I'm not familiar enough with the internals to know whether this is happening in JAX or somewhere lower-level. However, here's an attempted workaround, via a custom reverse rule.

```python
from functools import partial
from jax import custom_vjp

@partial(custom_vjp, nondiff_argnums=(2,3))
def sparse_matrix_product(target, values, nonzeroes, shape):
    return BCOO((values, nonzeroes), shape = shape, unique_indices = True) @ target

def sparse_matrix_product_fwd(target, values, nonzeroes, shape):
    out = sparse_matrix_product(target, values, nonzeroes, shape)
    carry = (target, values)
    return out, carry

def sparse_matrix_product_rev(nonzeroes, shape, carry, cotangents):
    (target, values) = carry
    target_cotangents = BCOO((values, nonzeroes), shape = shape, unique_indices = True).T @ cotangents

    x1 = nonzeroes[...,0]
    x2 = nonzeroes[...,1]

    sparse_cotangents = cotangents[...,x1,:]
    sparse_targets = target[...,x2,:]
    values_cotangents = jnp.sum(sparse_cotangents * sparse_targets, axis=-1)

    return (target_cotangents, values_cotangents)

sparse_matrix_product.defvjp(sparse_matrix_product_fwd, sparse_matrix_product_rev)



import jax 
import jax.numpy as jnp
from jax.experimental.sparse import BCOO

n = 1000
m = 777
p = 11
nonzeroes = jnp.stack((jnp.arange(0,m), jnp.arange(p,m+p)), axis=-1)

def quadratic(p):
    x1 = nonzeroes[...,0]
    x2 = nonzeroes[...,1]
    values = (x1 - x2)**2 / p
    matrix = BCOO((values, nonzeroes), shape=(n,n), unique_indices=True)
    return jnp.ones(n).T @ matrix @ jnp.ones(n)

def fixed_quadratic(p):
    x1 = nonzeroes[...,0]
    x2 = nonzeroes[...,1]
    values = (x1 - x2)**2 / p
    return (jnp.ones(n).T @ sparse_matrix_product(jnp.expand_dims(jnp.ones(n),-1), values, nonzeroes, (n,n))).squeeze()

(jax.grad(quadratic)(1.0)), jax.grad(fixed_quadratic)(1.0)

graph = jax.xla_computation(jax.grad(fixed_quadratic))(1.0)

with open("t.dot", "w") as f:
    f.write(graph.as_hlo_dot_graph())
```

Here's the corresponding computational graph: [t.pdf](https://github.com/google/jax/files/10806692/t.pdf). This appears so far to avoid dense matrix assembly, though I have not done any testing yet beyond an interactive notebook and thus there may be mistakes.
Very nice! Unfortunately it's not that easy if you have to handle all the possibilities offered by `dot_general`, which the BCOO code has to account for. We could probably special-case some situations and do something like this in simpler cases.

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -1086,11 +1086,43 @@ def bcoo_dot_general_sampled(A: Array, B: Array, indices: Array, *, dimension_nu
   return bcoo_dot_general_sampled_p.bind(A, B, indices,
                                          dimension_numbers=(cdims, bdims))
 
+def _bcoo_dot_general_sampled_slow(A, B, indices, *, dimension_numbers):
+  return _bcoo_extract(indices, lax.dot_general(A, B, dimension_numbers=dimension_numbers))
+
+def _bcoo_dot_general_sampled_simple(A, B, indices, *, dimension_numbers):
+  (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers
+  assert not (lhs_contract or rhs_contract or lhs_batch or rhs_batch)
+  assert A.ndim == B.ndim == 1
+  n_batch = indices.ndim - 2
+  n_sparse = indices.shape[-1]
+  nse = indices.shape[-2]
+  assert n_batch + n_sparse == 2
+  if n_batch == 0:
+    return A[indices[:, 0]] * B[indices[:, 1]]
+  elif n_batch == 1:
+    return A[:, None] * B[indices[..., 0]]
+  elif n_batch == 2:
+    out = A[:, None, None] * B[None, :, None]
+    return lax.broadcast_in_dim(out, (len(A), len(B), nse), (0, 1, 2))
+  else:
+    raise ValueError("too many batch dimensions.")
+
 @bcoo_dot_general_sampled_p.def_impl
 def _bcoo_dot_general_sampled_impl(A, B, indices, *, dimension_numbers):
-  # TODO(jakevdp): use a more efficient implementation that avoids the full dot product.
-  dense_result = lax.dot_general(A, B, dimension_numbers=dimension_numbers)
-  return _bcoo_extract(indices, dense_result)
+  A = jnp.asarray(A)
+  B = jnp.asarray(B)
+  indices = jnp.asarray(indices)
+  (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers
+  n_batch = indices.ndim - 2
+  n_sparse = indices.shape[-1]
+
+  # TODO(jakevdp): add fast approach for more general cases.
+  if (not (lhs_contract or rhs_contract or lhs_batch or rhs_batch)
+      and A.ndim == B.ndim == 1 and n_sparse + n_batch == 2):
+    return _bcoo_dot_general_sampled_simple(A, B, indices, dimension_numbers=dimension_numbers)
+
+  return _bcoo_dot_general_sampled_slow(A, B, indices, dimension_numbers=dimension_numbers)
+
 
 @bcoo_dot_general_sampled_p.def_abstract_eval
 def _bcoo_dot_general_sampled_abstract_eval(A, B, indices, *, dimension_numbers):


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1294,6 +1294,30 @@ def sparse_fun(lhs, rhs, indices):
       # TODO(jakevdp) fix forward-mode autodiff & enable tests here.
       self._CheckGradsSparse(dense_fun, sparse_fun, args_maker, modes=['rev'], argnums=[0, 1])
 
+  @jtu.sample_product(
+    xshape=[(3,), (5,)],
+    yshape=[(3,), (5,)],
+    dtype=jtu.dtypes.floating + jtu.dtypes.complex,
+    n_batch=[0, 1, 2],
+  )
+  def test_bcoo_dot_general_sampled_fast(self, xshape, yshape, n_batch, dtype):
+    rng = jtu.rand_default(self.rng())
+    sprng = sptu.rand_bcoo(self.rng(), n_batch=n_batch)
+
+    dimension_numbers = (([], []), ([], []))
+    args_maker = lambda: [rng(xshape, dtype), rng(yshape, dtype),
+                          sprng(xshape + yshape, dtype).indices]
+
+    def f1(x, y, indices):
+      mat_full = lax.dot_general(x, y, dimension_numbers=dimension_numbers)
+      return sparse_bcoo._bcoo_extract(indices, mat_full)
+
+    def f2(x, y, indices):
+      return sparse.bcoo_dot_general_sampled(x, y, indices, dimension_numbers=dimension_numbers)
+
+    self._CheckAgainstNumpy(f1, f2, args_maker)
+    self._CompileAndCheck(f2, args_maker)
+
   @jtu.sample_product(
     [dict(n_batch=n_batch, n_dense=n_dense, lhs_shape=lhs_shape,
           rhs_shape=rhs_shape, dimension_numbers=dimension_numbers)