algorithmnumpymatrix-multiplicationnumpy-einsum

Product of arrays (einsum) of 3D arrays containing only -1 or +1


Let X be an array of shape (M, k, g) and Q be an array of shape (m, k, g), where m, M, k, and g can be "very large". Suppose the entries of X and Y are either -1 or plus +1. I'm interested in the array Z of shape (m, M, k) defined by Z[a,b,i] = Q[a,i] @ X[b,i]. It is clear that this can be accomplished in Numpy (or Jax, for GPU exploitation) like so

Z = einsum("mkg,Mkg->mMk", Q, X)

Question. Can Z be computed more efficiently by using the information that the entries of X and Y are -1 or +1 ?


Solution

  • As I mentioned in my comment, casting to int8 seems to give a speed up. If you are able and willing to use torch, I see a speedup when using their einsum (cannot elaborate on which optimizations they are using), seems slower if casting to their int8 type though:

    enter image description here