pairwise dot product with einsum

pcross.py

#!/usr/bin/env python3
import numpy as np
from time import time


a = np.random.random((10000, 3))
b = np.random.random((10000, 3))

t0 = time()
c = np.array([np.cross(x, y) for x, y in zip(a, b)])
print(f'vanilla: {time() - t0:.4f}')

t0 = time()
d = np.cross(a, b)

print(f'einsim: {time() - t0:.4f}')
print(f'identical? {np.isclose(c, d).all()}')

pdot.py

#!/usr/bin/env python3
import numpy as np
from time import time


a = np.random.random((1000000, 3))
b = np.random.random((1000000, 3))

t0 = time()
c = np.array([x@y for x,y in zip(a, b)])
print(f'vanilla: {time() - t0:.4f}')

t0 = time()
d = np.einsum('ij,ij->i', a, b)

print(f'einsim: {time() - t0:.4f}')
print(f'identical? {np.isclose(c, d).all()}')

pdot_2d.py

import numpy as np
from time import time


a = np.random.random((100000, 3, 3))
b = np.random.random((100000, 3, 3))


# traditional pairwise dot product
t0 = time()
d = np.array([x @ y for x, y in zip(a, b)])
print(f"vanilla: {(time() - t0) * 1e3:.2f}ms")

# use the Einstein summation
t0 = time()
c = np.einsum('ijk,ikl->ijl', a, b)
print(f"einsum: {(time() - t0) * 1e3:.2f} ms")

print("identical?", np.allclose(c, d))

The result that I got is

(pdot)
vanilla: 0.9311
einsim: 0.0060
identical? True

(pcross)
vanilla: 0.2328
einsim: 0.0002
identical? True

(pdot2d)
vanilla: 253.31ms
einsum: 55.31 ms
identical? True

einsum is much faster than doing dot products inside a for loop. Maybe there is a better way but I don't know.