Broadcasting in Python with permutations
<p>I understand that <code>transpose</code> on an <code>ndarray</code> is intended to be the equivalent of matlab's <code>permute</code> function however I have a specific usecase that doesn't work simply. In matlab I have the following:</p>

<pre><code>C = @bsxfun(@times, permute(A,[4,2,5,1,3]), permute(B, [1,6,2,7,3,4,5])
</code></pre>

<p>where A is a 3D tensor of shape NxNxM and B is a 5D tensor of shape NxNxMxPxP. The above function is meant to vectorize looped kronecker products. I'm assuming that Matlab is adding 2 singleton dimensions for both A and B which is why it's able to rearrange them. I'm looking to port this code over to Python <strike>but I don't think it has the capability of adding these extra dimensions.</strike>. I found <a href="http://stackoverflow.com/questions/17835121/is-this-the-best-way-to-add-an-extra-dimension-to-a-numpy-array-in-one-line-of-c">this</a> which successfully adds the extra dimensions however the broadcasting is not working the same matlab's <code>bsxfun</code>. I have attempted the obvious translation (yes I am using numpy for these <code>ndarray</code>'s and functions):</p>

<pre><code>A = A[...,None,None]
B = B[...,None,None]
C = transpose(A,[3,1,4,0,2])*transpose(B,[0,5,1,6,2,3,4])
</code></pre>

<p><strike>and I get the following error:</p>

<pre><code>return transpose(axes)
ValueError: axes don't match array
</code></pre>

<p>My first guess is to do a <code>reshape</code> on A and B to add in those singleton dimensions?</strike></p>

<p>I now get the following error:</p>

<pre><code>mults = transpose(rho_in,[3,1,4,0,2])*transpose(proj,[0,5,1,6,2,3,4])
ValueError: operands could not be broadcast together with shapes (1,9,1,9,8) (9,1,9,1,8,40,40)
</code></pre>

<p>EDIT: Amended my question to be less about adding singleton dimensions but more about correctly broadcasting this matlab multiplication in python.</p>

<p>Looking at your MATLAB code, you have -</p>

<pre><code>C = bsxfun(@times, permute(A,[4,2,5,1,3]), permute(B, [1,6,2,7,3,4,5])
</code></pre>

<p>So, in essence -</p>

<pre><code>B : 1 , 6 , 2 , 7 , 3 , 4 , 5 
A : 4 , 2 , 5 , 1 , 3
</code></pre>

<p>Now, in MATLAB we had to borrow singleton dimensions from higher ones, that's why all that trouble of bringing in dims <code>6</code>, <code>7</code> for <code>B</code> and dims <code>4</code> <code>5</code> for <code>A</code>.</p>

<p>In NumPy, we bring in those explicitly with <a href="http://docs.scipy.org/doc/numpy-1.10.0/reference/arrays.indexing.html#numpy.newaxis" rel="nofollow"><code>np.newaxis</code>/None</a>. Thus, for NumPy, we could put it like so -</p>

<pre><code>B : 1 , N , 2 , N , 3 , 4 , 5 
A : N , 2 , N , 1 , 3 , N , N
</code></pre>

<p>, where <code>N</code> represents new axis. Please notice that we were need to put in new axes at the end  for <code>A</code> to push forward other dimension for alignment. In constrast, this happens in MATLAB by default.</p>

<p>Making that <code>B</code> looks easy enough with as the dimensions seem to be in order and we just need to add in new axes at appropriate places - <code>B[:,None,:,None,:,:,:]</code>.</p>

<p>Creating such <code>A</code> doesn't look straight-forward. Ignoring <code>N's</code> in <code>A</code>, we would have - <code>A : 2 , 1 , 3</code>. So, the starting point would be permuting dimensions and then add in those two new axes that were ignored - <code>A.transpose(1,0,2)[None,:,None,:,:,None,None]</code>.</p>

<p>Thus far, we have -</p>

<pre><code>B (new): B[:,None,:,None,:,:,:]
A (new): A.transpose(1,0,2)[None,:,None,:,:,None,None]
</code></pre>

<p>In NumPy, we can skip the leading new axes and trailing non-singleton dims. So, we could simplify like so -</p>

<pre><code>B (new): B[:,None,:,None]
A (new): A.transpose(1,0,2)[:,None,...,None,None]
</code></pre>

<p>The final output would be multiplication between these two extended versions -</p>

<pre><code>C = A.transpose(1,0,2)[:,None,...,None,None]*B[:,None,:,None]
</code></pre>

<p><strong>Runtime test</strong></p>

<p>I believe @Andras's post meant the equivalent <code>np.einsum</code> implementation to be something like : <code>np.einsum('ijk,lmkno-&gt;ljmikno',A,B)</code>.</p>

<pre><code>In [24]: A = np.random.randint(0,9,(10,10,10))
    ...: B = np.random.randint(0,9,(10,10,10,10,10))
    ...: 

In [25]: C1 = np.einsum('ijk,lmkno-&gt;ljmikno',A,B)

In [26]: C2 = A.transpose(1,0,2)[:,None,...,None,None]*B[:,None,:,None]

In [27]: np.allclose(C1,C2)
Out[27]: True

In [28]: %timeit np.einsum('ijk,lmkno-&gt;ljmikno',A,B)
10 loops, best of 3: 102 ms per loop

In [29]: %timeit A.transpose(1,0,2)[:,None,...,None,None]*B[:,None,:,None]
10 loops, best of 3: 78.4 ms per loop

In [30]: A = np.random.randint(0,9,(15,15,15))
    ...: B = np.random.randint(0,9,(15,15,15,15,15))
    ...: 

In [31]: %timeit np.einsum('ijk,lmkno-&gt;ljmikno',A,B)
1 loop, best of 3: 1.76 s per loop

In [32]: %timeit A.transpose(1,0,2)[:,None,...,None,None]*B[:,None,:,None]
1 loop, best of 3: 1.36 s per loop
</code></pre>

python matlab vectorization broadcasting permute