Relax assumptions for bcoo_dot_general

[ryan112358]

I'm trying to do some einsums with some sparse matrices and am hitting an error. Here is a minimal example:

def foo(ab, a):
  return jnp.einsum('ab,a->a', ab, a)

A = jax.random.normal(jax.random.key(0), shape=(2,))
AB = jax.random.normal(jax.random.key(1), shape=(2, 5))
AB = sparse.BCOO.fromdense(AB)

ans = jax.sparsify(foo)(AB, A)
NotImplementedError: bcoo_dot_general batch dimensions must be among the batch dimensions in the sparse representation.
got lhs_batch=(0,), n_batch=0

Would be nice if this type of operatoin could be natively supported. In this simple case, one could just write something like: ab.data *= a[ab.indices[:, 0]] then do ab.sum(axis=1), or slightly better do ab.sum(axis=1) * a

[jakevdp]

Hi - thanks for the request! We don't have plans to further develop jax.experimental.sparse at this point, but we would consider pull requests in this area if you're willing to work on this feature. Thanks!

[Aniketsy]

Hi @jakevdp ,
I’m interested in working on this feature to support batch dimensions in bcoo_dot_general for einsum-like operations with sparse matrices. Before proceeding, I’d appreciate some guidance on whether this aligns with JAX’s future plans for sparse support and if there are any design constraints or preferred implementation approaches.

If this is a reasonable contribution, I’d be happy to submit a PR. Could you point me to relevant parts of the codebase or existing discussions that might help?

[ryan112358]

@Aniketsy personally this would be a very welcome contribution. This seems like a good place to start:

https://github.com/jax-ml/jax/blob/main/jax/experimental/sparse/bcoo.py#L681

[Aniketsy]

@ryan112358 Thank you for suggesting me . I will work on this.