Combining sparse and einsum to perform large sparse sum
I have a matrix A with shape=(N, N) and a matrix B with the same shape=(N, N). I am constructing a matrix M using the following einsum (using the opt_einsum library):
M = oe.contract('nm,in,jm,pn,qm->ijpq', A, B, B, B, B)
This is evaluating the following sum:
This yeilds matrix M with shape (N, N, N, N). I then reshape this to a 2D array of shape (N**2, N**2)
M = M.reshape((N**2, N**2))
This must be 2D as it is treated as a linear operator.
I want to use the sparse library, as M is sparse, and becomes too large to store for large N. I can make A and B sparse, and insert them into the oe.contract.
The problem is, sparse only supports 2D arrays and so fails to produce the 4D output of shape (N, N, N. N). Is there a way to combine the einsum and reshape steps to allow sparse to be used in this way, as the final shape of M is 2D?
I did not correctly interpret the error given by opt_einsum.
The problem is not that sparse does not support ND sparse arrays (it does!), but that I was not using a true einsum, as the indices summed over appear more than twice (n and m). As stated in the opt_einsum documentation this will result in the use of the sparse.einsum function, of which none exists. Using only 1 or 2 of each index works. Using a differerent path, one suggested for example by hpaulj can be used to solve the problem.