title: Einsum is slow for tensor multiplication
body: <p>I'm trying to optimize a particular piece of code to calculate the mahalanobis distance in a vectorized manner. I have a standard implementation which used traditional python multiplication, and another implementation which uses einsum. However, I'm surprised that the einsum implementation is slower than the standard python implementation. Is there anything I'm doing inefficiently in einsum, or are there potentially other methods such as tensordot that I should be looking into?</p>
  <pre><code>#SETUP
  BATCH_SZ = 128
  GAUSSIANS = 100
  
  xvals = np.random.random((BATCH_SZ, 1, 4))
  means = np.random.random((GAUSSIANS, 1, 4))
  inv_covs = np.random.random((GAUSSIANS, 4, 4))
  </code></pre>
  <pre><code>%%timeit
  xvals_newdim = xvals[:, np.newaxis, ...]
  means_newdim = means[np.newaxis, ...]
  diff_newdim = xvals_newdim - means_newdim
  regular = diff_newdim @ inv_covs @ (diff_newdim).transpose(0, 1, 3, 2)
  
  &gt;&gt; 731 µs ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
  </code></pre>
  <pre><code>%%timeit
  diff = xvals - means.squeeze(1)
  einsum = np.einsum(&quot;ijk,jkl,ijl-&gt;ij&quot;, diff, inv_covs, diff)
  
  &gt;&gt; 949 µs ± 22.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
  </code></pre>