Einstein sum notation

einsum.py

import numpy as np

rng = np.random.RandomState(0)

print "Trace"
A = rng.rand(3, 3)
print np.trace(A)
print np.einsum("ii", A)
print

print "Dot product"
a = rng.rand(5)
b = rng.rand(5)
print np.dot(a, b)
print np.einsum("i,i", a, b)
print

print "Outer product"
print np.outer(a, b)
print np.einsum("i,j", a, b)
print

print "Identity"
B = rng.rand(3, 2)
print B
print np.einsum("ij", B)
print

print "Transpose"
print B.T
print np.einsum("ji", B)
print

print "Sum"
print np.einsum('i->', a)
print np.sum(a)
print

print "Diag"
print np.einsum("ii->i", A)
print np.diag(A)
print

print "Matrix product"
print np.einsum("ij,jk", A, B)
print np.dot(A, B)
print