Sparsify and gradient calculations

[quattro]

Hi all, I've been prototyping a statistical model and wanted to make use of BCOO to reduce memory overhead and hopefully squeeze out some performance gains. So far the BCOO implementation has worked out well in reducing memory, but I've run into an issue when trying to solve a subroutine that performs some optimization.

The general framework does something like:

for loop in main_loop:
    params = _inner_loop(data, params)

@sparsify
def _inner_loop(data, params):
    params = # use jax grad to perform numerical optimization for some of the params
    params = # use some analytic results to update other params
    return params

There is a need to sparsify the entire inner loop to leverage BCOO supported operations throughout the code.

However, JAX complains that it cannot evaluate the gradient due to unsupported operations between add_any. I've coded up a reduced version of the problem to focus on the numerical aspect which highlights three different ways of when to attach the @sparsify attribute and when the problem arises here.

Running the linked gist results in the following output:

simulate data
init model
Attempting outer sparsify 1
Outer sparsify 1 threw exception sparse rule for add_any is not implemented.
Attempting wrapped sparsify 1
Wrapped sparsify 1 threw exception sum requires ndarray or scalar arguments, got <class 'jax.experimental.sparse.bcoo.BCOO'> at position 0.
Attempting wrapped sparsify 2
Wrapped sparsify 2 worked

I would greatly appreciate any help or comments on the matter.

[jakevdp]

Thanks for the question! jax.experimental.sparse is under active development, and we typically implement new primitives as they are needed. You are the first user I'm aware of to need the sparsification rule for add_any. Would you be able to put together a minimal example of code that requires this path?

[quattro]

Thanks @jakevdp . My confusion was along the lines of how different the behavior of the code was based on where sparsify was added. I will try to pare the gist above down into a minimal/readable example that generates the same issues.

I wouldn't necessarily need an add_any, but it seems to crop up during the gradient over jnp.sum((x ** 2) * wgt) where x is a sparse matrix and wgt is an intermediate result that depends on the parameter.

[jakevdp]

I see, thanks. Keep in mind that sparsify(grad(f))(x) and grad(sparsify(f))(x) are, by design, fundamentally different computations. The former will differentiate with respect to all elements in the array, the latter will differentiate with respect to only the nonzero elements. The behavior makes sense if you think about how the transformations compose.

[quattro]

Wonderful explanation, thanks so much. That seems obvious in hindsight but wasn't something I had appreciated before. Cheers and thanks again.