Improve efficiency of sparse symmetric ERI

[blazejba]

Problem:

This blog post describes how to implement sparse matrix multiplication in Jax as:

@jax.partial(jax.jit, static_argnums=(2))
def sp_matmul(A, B, shape):
    """
    Arguments:
        A: (N, M) sparse matrix represented as a tuple (indexes, values)
        B: (M,K) dense matrix
        shape: value of N
    Returns:
        (N, K) dense matrix
    """
    assert B.ndim == 2
    indexes, values = A
    rows, cols = indexes
    in_ = B.take(cols, axis=0)
    prod = in_*values[:, None]
    res = jax.ops.segment_sum(prod, rows, shape)
    return res

This has been done for ERI in sparse_symmetric_ERI.py in sequentialized_iter(...)

Alex generated a profile for it, and thinks it is not as efficient as it should

Solution:
Identify the source of inefficiency in sequentialized_iter(...) and find a work around.

[blazejba]

Alex identified that the main problem is actually in eigh, but there still might be oppties to speed up sequentialized_iter(...)