A count_nonzero function that's sparsify friendly

[BBischof]

In short:

I'm wondering if there's a version of count_nonzero that is sparsify friendly, an opportunity to add ne to sparse rules, or a good way to do this that could work for dense and sparse matrices?

More discussion

Hi, I'm in sparse-land again. (from jax.experimental import sparse) and today I was interested in calculating a mean with positive support, i.e. I want to know the mean of non-zero entries.

I'd really like the method to be friendly for both dense and sparse matrices (i.e. i'd like to write the function such that sparsify can manage the sparse version).

from jax.experimental import sparse
import jax.numpy as jnp
import numpy as np
from jax import grad, jit

a = jnp.array([[0., 1., 0.],
                [3., 0., 0.],
                [0., 0., 4.]])

def mean_with_positive_support(A):
  return jnp.sum(A)/jnp.count_nonzero(A)

mean_with_positive_support(a) # returns desired "DeviceArray(2.6666667, dtype=float32)"

sp_mean_with_positive_support = sparse.sparsify(mean_with_positive_support)

sp_mean_with_positive_support(
    sparse.BCOO.fromdense(a)
)

Which predictably returns:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-22-c572c43a75cc> in <module>()
      1 sp_mean_with_positive_support(
----> 2     sparse.BCOO.fromdense(a)
      3 )

10 frames
/usr/local/lib/python3.7/dist-packages/jax/experimental/sparse/transform.py in eval_sparse(jaxpr, consts, argspecs, spenv)
    309     if any(val.is_sparse() for val in invals):
    310       if prim not in sparse_rules:
--> 311         raise NotImplementedError(f"sparse rule for {prim}")
    312       out = sparse_rules[prim](spenv, *invals, **eqn.params)
    313     else:

NotImplementedError: sparse rule for ne

[jakevdp]

This is a great question, and I can't immediately think of any good answer. I'm currently exploring a similar issue here: #8991

In this case, it's a function for applying an operator element-wise. It sounds liike we could use a similar function that's something like sparse.reduce_nonzero(fun, x), that would basically return fun(x[x != 0]) in the case of a dense array, and fun(x.data) in case of a sparse array.

It would get a little tricky in the case of sparse representations with duplicate indices or with padding... we'd have to think about how to best do that.

[choltz95]

Also wanted to add another datapoint. I really appreciate the support for sparse arrays.
Have been using this implementation of sparse diag:

def sp_diag(M):
  m, n = M.shape
  diag = jnp.zeros(m)
  diag_idx = M_sp.indices[:,0]==M_sp.indices[:,1]
  diag = diag.at[M_sp.indices[diag_idx]].set(M_sp.data[diag_idx])
  return diag

but it's very brittle & prob not the right way to do things.

[jakevdp]

Thanks for the question. Here's one approach that's a bit more robust:

>>> M = sparse.BCOO.fromdense(jnp.arange(16).reshape((4, 4)))                             

>>> print(M.todense())                                                                    
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

>>> diag = (M * jnp.eye(*M.shape)).sum(1)                                                 

>>> print(diag.todense())                                                                 
[ 0.  5. 10. 15.]

This is roughly the way that jnp.diagonal is implemented right now, believe it or not (though it uses a selection instead of a multiplication to generate the sparse diagonal matrix)