Repository: google/jax

Issue:
Sparse rule for div is not implemented
### Description

Hi all, I've been prototyping some functionality using the experimental `BCOO` interface and quickly run into `NotImplementedError` for basic functionality. For example, computing the mean across a given axis. 

```python
import jax.experimental.sparse as xsp
import jax.numpy as jnp

X = jnp.hstack([jnp.ones(5)[:,jnp.newaxis], jnp.zeros(5)[:,jnp.newaxis]])
sX = xsp.BCOO.fromdense(X)
smean = xsp.sparsify(jnp.mean)
smean(sX)
```
Will output `Array(0.5, dtype=float32)`, however, `smean(sX, axis=0)` (or `axis=1`) throws 
```python
NotImplementedError: sparse rule for div is not implemented.
```

Is this expected at this time, and are there plans to extend the functionality for `div`?

### What jax/jaxlib version are you using?

jax v0.4.5

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.9.12

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the question - we haven't implemented sparse div because in general it's unsafe: if the divisor is zero, then it would be a densifying operation (all unspecified entries would become `NaN`).

Perhaps you could call `sum()` and then densify the result and divide by `N` to compute the mean?

Resolving patch:
diff --git a/jax/experimental/sparse/transform.py b/jax/experimental/sparse/transform.py
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@@ -671,6 +671,22 @@ def _mul_sparse(spenv, *spvalues):
 
 sparse_rules_bcoo[lax.mul_p] = _mul_sparse
 
+def _div_sparse(spenv, *spvalues):
+  X, Y = spvalues
+  if Y.is_sparse():
+    raise NotImplementedError(
+      "Division by a sparse array is not implemented because it "
+      "would result in dense output. If this is your intent, use "
+      "sparse.todense() to convert your arguments to a dense array.")
+  X_promoted = spvalues_to_arrays(spenv, X)
+  out_data = bcoo_multiply_dense(X_promoted, 1. / spenv.data(Y))
+  out_spvalue = spenv.sparse(X.shape, out_data, indices_ref=X.indices_ref,
+                              indices_sorted=X.indices_sorted,
+                              unique_indices=X.unique_indices)
+  return (out_spvalue,)
+
+sparse_rules_bcoo[lax.div_p] = _div_sparse
+
 def _reduce_sum_sparse(spenv, *spvalues, axes):
   X, = spvalues
   X_promoted = spvalues_to_arrays(spenv, X)
@@ -894,6 +910,8 @@ def _sparse_iter(arr):
   "__rmatmul__": sparsify(_swap_args(jnp.matmul)),
   "__mul__": sparsify(jnp.multiply),
   "__rmul__": sparsify(_swap_args(jnp.multiply)),
+  "__truediv__": sparsify(jnp.divide),
+  "__rtruediv__": sparsify(_swap_args(jnp.divide)),
   "__add__": sparsify(jnp.add),
   "__radd__": sparsify(_swap_args(jnp.add)),
   "__sub__": sparsify(jnp.subtract),


Test patch:
diff --git a/tests/sparsify_test.py b/tests/sparsify_test.py
--- a/tests/sparsify_test.py
+++ b/tests/sparsify_test.py
@@ -240,6 +240,33 @@ def testSparseMul(self, shape, dtype, n_batch, n_dense, unique_indices):
 
     self.assertAllClose(out.todense(), x.todense() * y.todense())
 
+  @jtu.sample_product(
+    [dict(shape=shape, n_batch=n_batch, n_dense=n_dense)
+      for shape in [(5,), (5, 8), (8, 5), (3, 4, 5), (3, 4, 3, 2)]
+      for n_batch in range(len(shape) + 1)
+      for n_dense in range(len(shape) + 1 - n_batch)
+    ],
+    dtype=jtu.dtypes.integer + jtu.dtypes.floating + jtu.dtypes.complex,
+  )
+  def testSparseDiv(self, shape, dtype, n_batch, n_dense):
+    rng_dense = jtu.rand_nonzero(self.rng())
+    rng_sparse = rand_sparse(self.rng(), rand_method=jtu.rand_some_zero)
+    x = BCOO.fromdense(rng_sparse(shape, dtype), n_batch=n_batch,
+                       n_dense=n_dense)
+    spdiv = self.sparsify(operator.truediv)
+
+    # Scalar division
+    divisor = 2
+    expected = x.todense() / divisor
+    self.assertAllClose(expected, spdiv(x, divisor).todense())
+    self.assertAllClose(expected,  (x / divisor).todense())
+
+    # Array division
+    divisor = rng_dense(shape, dtype)
+    expected = x.todense() / divisor
+    self.assertAllClose(expected, spdiv(x, divisor).todense())
+    self.assertAllClose(expected,  (x / divisor).todense())
+
   def testSparseSubtract(self):
     x = BCOO.fromdense(3 * jnp.arange(5))
     y = BCOO.fromdense(jnp.arange(5))