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

Question

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 ?

Answer 1

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:

---

Answer 2

Here are some more benchmarks on GPU (T4) to complement the accepted answer. I compared Numpy vs JAX vs PyTorch

CPU

GPU

The conclusion (based on this tiny experiment) is:

• Only Numpy knows how to take full advantage int8. The others are probably doing some weird casting under the hood.
• GPU helps JAX.
• PyTorch is a beast on GPU.