numpy einsum for multiplication looped along axis

Question

I have an 2 x 2 matrix yy

yy = np.array([[0.5, 0], [0, 2]])
print(yy)

array([[0.5, 0. ],
       [0. , 2. ]])

and n=3 x 4 x 2 matrix xy

xy = np.array([
    [[1, 0.1], [2, 0.2], [3, 0.3], [4, 0.4]],
    [[2, 0.2], [4, 0.3], [6, 0.4], [8, 0.5]],
    [[1, 0.3], [3, 0.4], [5, 0.5], [7, 0.6]],
])
print(xy)

array([[[1. , 0.1],
        [2. , 0.2],
        [3. , 0.3],
        [4. , 0.4]],

       [[2. , 0.2],
        [4. , 0.3],
        [6. , 0.4],
        [8. , 0.5]],

       [[1. , 0.3],
        [3. , 0.4],
        [5. , 0.5],
        [7. , 0.6]]])

I want to multiply yy on transposed elements of xy along 0 axis via numpy.einsum without loops to get

result = np.concatenate((np.expand_dims((yy @ xy[0].T).T, 0), np.expand_dims((yy @ xy[1].T).T, 0), np.expand_dims((yy @ xy[2].T).T, 0)), axis=0)
print(result)

array([[[0.5, 0.2],
        [1. , 0.4],
        [1.5, 0.6],
        [2. , 0.8]],

       [[1. , 0.4],
        [2. , 0.6],
        [3. , 0.8],
        [4. , 1. ]],

       [[0.5, 0.6],
        [1.5, 0.8],
        [2.5, 1. ],
        [3.5, 1.2]]])

that is every i-th element should be

i = 0
print((yy @ xy[i,:,:].T).T)

[[0.5 0.2]
 [1.  0.4]
 [1.5 0.6]
 [2.  0.8]]

How to achieve that with numpy.einsum?

Answer 1

or rewriting result:

 result1 = np.stack([(yy @ xyi.T).T for xyi in xy])

the equivalent einsum:

In [32]: np.allclose(np.einsum('ij,kmj->kmi', yy, xy),result)
Out[32]: True

And since the k dimension of xy is a 'batch' one, we can stick with matmul:

In [34]: np.allclose(xy@yy.T,result)
Out[34]: True

This would have been more obvious if we wrote the iterative version with fewer transpose:

result2 = np.stack([xyi@yy.T for xyi in xy])