4 / 6 NumPy & JAX

Batched Einsum

Let A have shape (8, 3, 5) representing (batch, n, m) and B have shape (8, 5, 7) representing (batch, m, p). Which np.einsum call computes the batched matrix product with output shape (8, 3, 7)?

A

np.einsum('bmn,bmp->bnp', A, B)

B

np.einsum('nm,mp->bnp', A, B)

C

np.einsum('bnm,bmp->bnp', A, B)

D

np.einsum('bnm,bpm->bnp', A, B)

E

np.einsum('bnm,bmp->pnb', A, B)