Using einsum without explicit replication

2
votes

I have a matrix A of shape (n, n) and another matrix b of shape (p, n). I need to get a matrix C such that 

C[i] = (A * b[i, np.newaxis, :]) * b[i, :, np.newaxis]

I'm currently doing that by doing np.matlib.repmat(A) and then using np.einsum as follows

A1 = np.matlib.repmat(A, p, 1).reshape(p, n, n)
C = np.einsum('ijk, ij..., ik... -> ijk', A1, b, b)

But I can see that creating A1 is a waste of resources as it's just using the same values over and over. Is there anyway I can do this without creating the intermediate matrix A1?

1
pythonarraysnumpynumpy-einsum

1 Answers

2
votes

Use np.einsum like so -

np.einsum('ij,kij,kij->kij',A,b[:,None,:],b[...,None])