Surprising results from squaring a sparse matrix

[BBischof]

I'll start by acknowledging that this is part of experimental and so I recognize that things are still in prototype.

I'll also acknowledge that this was my first attempt at playing with sparsify so I may be just mis-interpreting how things should work.

Here's what I'm seeing that came as a surprise to me:

I was expecting that the sparse matrix squaring would return the same result.

Here's the code to reproduce:

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.]])

u = jnp.array([[0., 4.],[0., 2.],[0., 0.],])
v = jnp.array([[2., 1.],[2., 4.],[0., 5.],])

def s_mult_proto(x,y):
  return jnp.multiply(x,y)

s_mult = sparse.sparsify(s_mult_proto)

m = a - u @ v.T
print(m)
s_mult_proto(m,m)

m = (sparse.BCOO.fromdense(a) - (sparse.BCOO.fromdense(u) @ sparse.BCOO.fromdense(v).T))
print(m.todense())
s_mult(m,m).todense()

[jakevdp]

Thanks for the report - this looks like a bug with the newly-added sparse lowering of lax.mul.

For the moment, I think you can work around it by doing

m = m.sum_duplicates()

after which s_mult(m, m).todense() will display the correct result.

I've opened #8889 to track the bug. Thanks!

[BBischof]

Hm, I may not understand the workaround:

m = (sparse.BCOO.fromdense(a) - (sparse.BCOO.fromdense(u) @ sparse.BCOO.fromdense(v).T))
print(m.todense())
m = m.sum_duplicates()
s_mult(m,m).todense()

Yields:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-5-c39f20f1bbc6> in <module>()
----> 1 m = (sparse.BCOO.fromdense(a) - (sparse.BCOO.fromdense(u) @ sparse.BCOO.fromdense(v).T))
      2 print(m.todense())
      3 m = m.sum_duplicates()
      4 s_mult(m,m).todense()

5 frames
/usr/local/lib/python3.7/dist-packages/jax/experimental/sparse/bcoo.py in _bcoo_spdot_general_abstract_eval(lhs_data, lhs_indices, rhs_data, rhs_indices, lhs_shape, rhs_shape, dimension_numbers)
    767 
    768   if not (rhs.n_sparse == 1):
--> 769     raise NotImplementedError("bcoo_spdot_general with n_sparse != 1 on the rhs.")
    770 
    771   if max(lhs_batch, default=-1) >= lhs.n_batch or max(rhs_batch, default=-1) >= rhs.n_batch:

NotImplementedError: bcoo_spdot_general with n_sparse != 1 on the rhs.

[jakevdp]

That error is being raised before the updated line. Most of this functionality is very new (only available in master, or perhaps in the most recent release). Are you sure you haven't changed to an older JAX version?

[BBischof]

oh, good eye. Apologies; lemme see what happened.

[BBischof]

Embarrassing mistake 😅 I copy pasted the code from my question and forgot to update JAX in the colab.

This suggestion worked:

[jakevdp]

Great! I'll probably dig-in tomorrow to see if I can fix the issue. Thanks for kicking the tires a bit on this code - please do let me know if you run into any other issues, or if you have feature requests (the development right now is mainly prioritized based on what features people ask me for 😁 )

[BBischof]

One followup I have is that I'm unsure how I'd integrate the suggested trick into a function I want to be able to support both sparse and not sparse operations:

def mse_observed_loss(A, U, V):
    estimator = -(U @ V.T)
    estimator = estimator.sum_duplicates()
    return jnp.mean(jnp.multiply(A + estimator, A + estimator))

Which I want to use with:

mse_observed_loss_sp = sparse.sparsify(mse_observed_loss)

mse_observed_loss_sp(
    sparse.BCOO.fromdense(a),
    sparse.BCOO.fromdense(u),
    sparse.BCOO.fromdense(v),
).todense()

Apologies if i'm being...... dense. (😉)

[jakevdp]

There's not really any good way to do that at the moment (we could create a primitive similar to sparse.todense() for this purpose, but haven't yet). In any case, the fix is in #8897 so hopefully this won't block you for too long.

[BBischof]

Awesome; thanks for the super quick resolution.

To ensure I'm clear: after that is merged, I'll be able to remove the sum_duplicates() line, and this will be able to work with sparsify in the way I was expecting

[jakevdp]

Yes.

[BBischof]

TYSM!