What does np.einsum('ij,jk->ik', A, B)
mean?
Basically, we have two input matrices A and B and an output matrix C is defined such that:
Cik=∑j=1nAijBjk
Note that this is the formula for standard matrix multplication.
- Repeated letters on either side of the
,
means values along those axes are multiplied together.
- Letters on the left of
->
that aren’t on the right of ->
means that values along that axis are summed together.
Examples
Vectors
Assuming we have vectors A,B∈Rn , then…
np.einsum('i', A)
=> returns a vector C∈Rn identical to A
Ci=Ai
np.einsum('i->', A)
=> sums all elements of A into a scalar C∈R
C=∑i=1nAi
np.einsum('i,i->i', A, B)
=> element wise multplication of A and B into a vector C∈Rn
Ci=∑i=1nAiBi
np.einsum('i,i', A, B)
=> dot product of A and B into a scalar C∈R, i.e. C=ATB
C=∑i=1nAiBi
np.einsum('i,j->ij', A, B)
=> outer product of A and B into a 2-dimensional matrix C∈Rn×n, i.e. C=A⊗B=ABT
Cij=AiBj
Matrices
Assuming we have matrices A∈Rn×m and B∈Rm×p , where m and p are compatible with n and m where appropriate, then…
np.einsum('ij', A)
=> returns a matrix C∈Rn×m identical to A
Cij=Aij
np.einsum('ji', A)
=> a matrix C∈Rm×n identical to AT
Cij=Aji
np.einsum('ij->', A)
=> sums all elements of A into scalar C∈R
C=∑i=1n∑j=1mAij
np.einsum('ii->i', A)
=> returns the diagonal elements of A as a vector C∈Rn, i.e. C=diag(A)
Ci=Aii
np.einsum('ii', A)
=> returns the trace of A as a scalar C∈R, i.e. C=trace(A)
C=∑i=1nAii
np.einsum('ij->i', A)
=> sums rows of A into a vector C∈Rn, i.e. C=np.sum(A,axis=1)
Ci=∑j=1mAij
np.einsum('ij->j', A)
=> sums columns of A into a vector C∈Rm, i.e. C=np.sum(A,axis=0)
Cj=∑i=1nAij
np.einsum('ij,ij->ij', A, B)
=> elementwise product of A∈Rn×m and B∈Rn×m into matrix C∈Rn×m, i.e. C=A⊙B
Cij=AijBij
np.einsum('ij,jk->ik', A, B)
=> matrix multplication of A∈Rn×m and B∈Rm×p into matrix C∈Rn×p, i.e. C=AB
Cik=∑j=1mAijBjk
References