4

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:

enter image description here

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?

Jack Wetherell
  • 565
  • 3
  • 8
  • 16
  • 1
    `opt_einsum` is a third party library that few of us have worked with (is there even a `tag` for it?). A quick glance at the docs does not indicate whether it can work with `sparse` matrices. `np.einsum` certainly can't. – hpaulj Feb 13 '21 at 17:51
  • 1
    @hpaulj `opt_einsum` certainly does support `sparse` matrices, as indicated in the docs (https://optimized-einsum.readthedocs.io/en/stable/backends.html): "The sparse library also fits the requirements and is supported...." – Jack Wetherell Feb 13 '21 at 19:31
  • @hpaulj FWIW `opt_einsum` gets about 6M downloads/month, many use it without knowing it (tensorflow is a prime example) and the older algorithms are implemented in `np.einsum`. It doesn't always work as well as you would like in SO examples, as they are typically small meaning BLAS costs more than direct einsum and/or the scaling issues with naive einsum do not play a large role. (Disclaimer: author of opt_einsum) – Daniel Jan 15 '22 at 17:33

2 Answers2

1

This may not help with your use of opt_einsum, but with a bit of reorganizing I can speed up np.einsum quite a bit, at least for small arrays.

Do a partial product of two B:

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N,N)

The pq pair is the same, so we don't need to recalculate it:

c2 = np.einsum('nm,onm,pnm->op',A,c1,c1)

I verified that this works for two (3,3) arrays, and the speed up is about 10x.

We can even reshape the nm to 1d, though this doesn't improve speed:

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N*N)
c3 = np.einsum('n,on,pn->op',A.reshape(N*N),c1,c1)
hpaulj
  • 221,503
  • 14
  • 230
  • 353
1

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.

Jack Wetherell
  • 565
  • 3
  • 8
  • 16
  • 1
    Yes, in this case we need to cast the spare arrays back to dense NumPy ndarray which is quite suboptimal. We have a hard time finding general heuristics of how to perform the above outer products to improve GEMM performance, but if you have ideas let us know here: https://github.com/dgasmith/opt_einsum – Daniel Jan 15 '22 at 17:35