Search code examples
pythonnumpytensortensordot

Replicating Matrix-Matrix Multiplication using Tensordot


In PyKeops package, there is no available formula for Matrix-Matrix multiplication. Instead, they have implemented something similar to numpy.tensordot. I have two matrices A,B of size m x n and n x n. Is there any way to replicate A @ B using numpy.tensordot.

import numpy as np
m,n = 10, 20
A = np.random.random((m,n))
B = np.random.random((n,n))

result1 = A @ B
result2 = np.tensordot(A,B,(1,1))

print(result1 == result2)

# =======================
# array([[False, False, False, False, False, False, False, False, 
# ....

Solution

  • Okay I found an answer to this, we can fix this using the following parameters given to tensordot:

    import numpy as np
    m,n = 10, 20
    A = np.random.random((m,n))
    B = np.random.random((n,n))
    
    result1 = A @ B
    result2 = np.tensordot(A,B,[(1,),(0,)])
    
    print(result1 == result2)
    
    # =======================
    # array([[True,  True,  True,  True,  True,  True,  True,  True,  True, 
    # ....