title: In Multi-dimensional array products, how do I align axes with and without summation?
body: <p>What is the best way to do array operations when there are some repeated indices which are summed over AND others which are not?  It seems like I may have to use <code>einsum</code> for these operations, but it would be better if there was a <code>tensordot</code> alternative with a flag for dimensions that are aligned but not summed.  </p>
  
  <p>Does anybody know of a fast numerical routine (in lapack perhaps?) that behaves like tensordot, except that some axes can be aligned without being summed?</p>
  
  <p>==</p>
  
  <p>Here is an example code to show type of array operation that is needed.  The operation I need is done by <code>method_sum</code>, <code>method_einsum</code>, and <code>method_matmul</code>.  A similar operation sums over the matching j-axis and is done by <code>method2_einsum</code> and <code>method2_tensordot</code>.  </p>
  
  <p>By comparing times it seems that <code>tensordot</code> should be able to beat <code>einsum</code> on the first problem.  However, it does not have the functionality to align axes without summing them up.     </p>
  
  <pre><code>#import scipy
  import scipy as sp
  
  # Shapes of arrays
  I = 200
  J = 50
  K = 200
  L = 100
  
  a = sp.ones((I, J, L))
  b = sp.ones((J, K, L))
  
  
  # The desired product has a sum over the l-axis
  
  ## Use broadcasting to multiply and sum over the last dimension
  def method_sum(a, b):
      "Multiply arrays and sum over last dimension."   
      c = (a[:, :, None, :] * b[None, :, :, :]).sum(-1)
      return c
  
  ## Use einsum to multiply arrays and sum over the l-axis
  def method_einsum(a, b):
      "Multiply arrays and sum over last dimension."
      c = sp.einsum('ijl,jkl-&gt;ijk', a, b)
      return c
  
  ## Use matmul to multiply arrays and sum over one of the axes
  def method_matmul(a, b):
      "Multiply arrays using the new matmul operation."
      c = sp.matmul(a[:, :, None, None, :], 
                    b[None, :, :, :, None])[:, :, :, 0, 0]
      return c
  
  
  # Compare einsum vs tensordot on summation over j and l
  
  ## Einsum takes about the same amount of time when j is not summed over) 
  def method2_einsum(a, b):
      "Multiply arrays and sum over last dimension."
      c = sp.einsum('ijl,jkl-&gt;ik', a, b)
      return c
  
  ## Tensor dot can do this faster but it always sums over the aligned axes
  def method2_tensordot(a, b):
      "Multiply and sum over all overlapping dimensions."
      c = sp.tensordot(a, b, axes=[(1, 2,), (0, 2,)])
      return c
  </code></pre>
  
  <p>Here are some times for the various routines on my computer.  Tensordot can beat einsum for <code>method2</code> because it uses multiple cores.  I would like to achieve performance like <code>tensordot</code> for the calculation where both the J and L-axes are aligned, but only the L-axis is summed. </p>
  
  <pre><code>Time for method_sum:
  1 loops, best of 3: 744 ms per loop
  
  Time for method_einsum:
  10 loops, best of 3: 95.1 ms per loop
  
  Time for method_matmul:
  10 loops, best of 3: 93.8 ms per loop
  
  Time for method2_einsum:
  10 loops, best of 3: 90.4 ms per loop
  
  Time for method2_tensordot:
  100 loops, best of 3: 10.9 ms per loop
  </code></pre>