Element-wise multiplication of BCOO sparse matrices reports out of memory error

[shirakawatoru]

I would like to use jax.experimental.sparse.bcoo_multiply_sparse:

import jax
import jax.experimental.sparse
X = jax.random.bernoulli(jax.random.PRNGKey(1234), 0.01, (5, 1000, 1000))
X_sp = jax.experimental.sparse.BCOO.fromdense(X)
jax.experimental.sparse.bcoo_multiply_sparse(X_sp, X_sp)

But this code reports the following error message:

---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-1-952fc59dc1ff> in <module>
      5 X_sp = jax.experimental.sparse.BCOO.fromdense(X)
----> 6 jax.experimental.sparse.bcoo_multiply_sparse(X_sp, X_sp)

19 frames
UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 10046854756 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:    2.34GiB
              constant allocation:         4B
        maybe_live_out allocation:    9.36GiB
     preallocated temp allocation:   10.53GiB
  preallocated temp fragmentation:       132B (0.00%)
                 total allocation:   22.22GiB
              total fragmentation:    1.17GiB (5.26%)
Peak buffers:

I think element-wise multiplication of COO matrices needs at most the memory size of original matrices. How can I fix this problem?

[jakevdp]

Hi - thanks for the question. You're right that in theory this should be possible: the problem is that jax.experimental.sparse is built on the operations available in XLA, and XLA doesn't have any set-like methods which are required for many sparse-sparse operations. So jax.experimental.sparse relies on a very inefficient method to compute sparse-sparse elementwise multiplication. I don't know of any workaround aside from using smaller matrices or a larger machine, or perhaps converting your matrices to a more structured sparse form (incidentally, issues like this are one reason JAX's sparse support is still experimental: it works well in many cases, but for sparse-sparse operations in particular it can be sub-optimal)

Hope that helps!

[shirakawatoru]

Thank you for answering my question. It helped me a lot. I understand that set operations cannot always have efficient linear representations. I will consider adopting a more structured sparse form for my problem.