user2498193
user2498193

Reputation: 1104

Cannot understand einsum calculation

I'm trying to port some code from python to R (forewarning.. I know far less python than I do R).

Anyway I've a complicated einsum command that I need to understand well in order to translate it and I'm having troubles. Based on an answer to previous question: https://stackoverflow.com/a/59858877/2498193 , I'm trying to code the calculation explicilty and having problems.

Reproducible example:

A = np.random.rand(2, 5, 1)
B = np.random.rand(2, 4, 5, 3)

C = np.einsum('ikl,ijkl->ijl', A, B)

""" Attempt to breakdown einsum into its parts"""
D = np.zeros((A.shape[0], B.shape[1], B.shape[3]))
for i in range(A.shape[0]): 
    for j in range(B.shape[1]):
        for l in range(B.shape[3]):
            for k in range(A.shape[1]):
                D[i, j, l] += A[i, k, l] * B[i, j, k, l]

So the calculation off C runs fine, but when I try D I get:

IndexError: index 1 is out of bounds for axis 2 with size 1

Where have I gone wrong?

Edit: I found an error in the indexing and fixed that but now have a new mystery

My fixed code for D and output:

for i in range(A.shape[0]): 
    for j in range(B.shape[1]):
        for l in range(A.shape[2]):
            for k in range(A.shape[1]):
                D[i, j, l] += A[i, k, l] * B[i, j, k, l]

C
Out[763]: 
array([[[1.12003067, 1.10913818, 0.6906052 ],
        [1.90393492, 0.95523242, 1.54739457],
        [0.94529917, 0.86866832, 1.50882582],
        [1.7744814 , 1.90689624, 1.55044583]],

       [[0.99766459, 0.81698549, 0.91783172],
        [1.05860284, 1.41360977, 0.84931137],
        [0.92352849, 1.40863288, 1.14309574],
        [0.95503075, 1.1450589 , 1.32160452]]])

D
Out[764]: 
array([[[1.12003067, 0.        , 0.        ],
        [1.90393492, 0.        , 0.        ],
        [0.94529917, 0.        , 0.        ],
        [1.7744814 , 0.        , 0.        ]],

       [[0.99766459, 0.        , 0.        ],
        [1.05860284, 0.        , 0.        ],
        [0.92352849, 0.        , 0.        ],
        [0.95503075, 0.        , 0.        ]]])

So I'm getting partially correct answer but unclear to me where all the zeroes are coming from ?

Upvotes: 0

Views: 78

Answers (1)

user2498193
user2498193

Reputation: 1104

I eventually figured out that I need to allow for broadcasting also in my breakdown of what einsum is doing.

Solution:

A = np.random.rand(2, 5, 1)
B = np.random.rand(2, 4, 5, 3)
C = np.einsum('ikl,ijkl->ijl', A, B)

D = np.zeros((A.shape[0], B.shape[1], B.shape[3]))
AA = np.zeros((2,5,3))
AA[:,:, 0] = A[:,:,0]
AA[:,:, 1] = A[:,:,0]
AA[:,:, 2] = A[:,:,0]

for i in range(AA.shape[0]): 
    for j in range(B.shape[1]):
        for l in range(AA.shape[2]):
            for k in range(AA.shape[1]):
                D[i, j, l] += AA[i, k, l] * B[i, j, k, l]

Out[946]: 
array([[[1.34883962, 1.29672134, 2.40826835],
        [2.12198906, 1.57248206, 1.62157716],
        [1.95668114, 1.45167364, 1.1761399 ],
        [2.14619827, 1.62883231, 1.42584051]],

       [[1.01555102, 1.48221712, 0.67038112],
        [0.52555659, 0.74096584, 0.73631941],
        [0.76738584, 0.61414461, 1.22202416],
        [0.5116698 , 0.94099001, 1.01196491]]])

D
Out[947]: 
array([[[1.34883962, 1.29672134, 2.40826835],
        [2.12198906, 1.57248206, 1.62157716],
        [1.95668114, 1.45167364, 1.1761399 ],
        [2.14619827, 1.62883231, 1.42584051]],

       [[1.01555102, 1.48221712, 0.67038112],
        [0.52555659, 0.74096584, 0.73631941],
        [0.76738584, 0.61414461, 1.22202416],
        [0.5116698 , 0.94099001, 1.01196491]]])

Shame r does not seem to have an numpy.einsum equivalent, but at least this is something I can implement!

Upvotes: 1

Related Questions