CODEWITHSUNDEEP
pythonnumpy
dzang
dzang

Reputation: 2260

What's the meaning of 'a' in np.einsum string?

I am trying to convert some code to run with Numba. np.einsum is not supported therefore I am trying to replace it with Numba supported functions.

I understood in part how np.einsum works and for example I get that:

x, y, z = 3, 2, 4
A = np.arange(x * y * z).reshape(x, y, z)
B = np.arange(x * y).reshape(x, y)

C = np.einsum('ijk,kj->ki', A.T, B)

is equivalent to:

C = np.sum(A.T * B.T, axis=1).T

e.g. I take ijk and the 3D canonical indices, but now I have the following expression which I cannot understand:

C = np.einsum('aij,jka->ajk', A, B)

What is the meaning of the index 'a' ? What would be an equivalent transformation using multiplication, sum and transposing?

Upvotes: 0

Views: 63

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53089

Which letters you use in the axes string doesn't matter much (but see bottom of this post), for example we can put z for a:

>>> A = np.arange(3*4*5).reshape(3,4,5)
>>> B = np.arange(5*2*3).reshape(5,2,3)
>>> 
>>> np.einsum('aij,jka->ajk',A,B)
array([[[   0,   90],
        [ 204,  306],
        [ 456,  570],
        [ 756,  882],
        [1104, 1242]],

       [[ 110,  440],
        [ 798, 1140],
        [1534, 1888],
        [2318, 2684],
        [3150, 3528]],

       [[ 380,  950],
        [1552, 2134],
        [2772, 3366],
        [4040, 4646],
        [5356, 5974]]])
>>> np.einsum('zij,jkz->zjk',A,B)
array([[[   0,   90],
        [ 204,  306],
        [ 456,  570],
        [ 756,  882],
        [1104, 1242]],

       [[ 110,  440],
        [ 798, 1140],
        [1534, 1888],
        [2318, 2684],
        [3150, 3528]],

       [[ 380,  950],
        [1552, 2134],
        [2772, 3366],
        [4040, 4646],
        [5356, 5974]]])

Equivalent without einsum:

>>> A.sum(1)[..., None]*B.transpose(2,0,1)
array([[[   0,   90],
        [ 204,  306],
        [ 456,  570],
        [ 756,  882],
        [1104, 1242]],

       [[ 110,  440],
        [ 798, 1140],
        [1534, 1888],
        [2318, 2684],
        [3150, 3528]],

       [[ 380,  950],
        [1552, 2134],
        [2772, 3366],
        [4040, 4646],
        [5356, 5974]]])

The identity of index letters matters where the output axes are implicit, because they are then assumed to be in alphabetical order

>>> A = np.ones((2,1))
>>> np.einsum('ab', A)
array([[1.],
       [1.]])
>>> np.einsum('zb', A)
array([[1., 1.]])

Upvotes: 3

Related Questions