将两个3d numpy数组合并为带有einsum或广播的2d数组？

Question

我有循环将两个3-d numpy数组中的切片相乘以生成2-d数组。我怀疑必须采用更有效的方法来进行广播或einsum，但我无法弄清楚。

A = np.array([[[ 9.99948603e-01, -9.99948603e-01,  0.00000000e+00,
                -0.00000000e+00,  0.00000000e+00, -0.00000000e+00,
                 0.00000000e+00, -0.00000000e+00],
               [ 9.91155487e-08, -3.34504020e-04,  0.00000000e+00,
                -0.00000000e+00,  9.91155487e-08, -3.34504020e-04,
                 0.00000000e+00, -0.00000000e+00]],

              [[ 8.80751720e-01, -9.97142430e-01,  0.00000000e+00,
                -0.00000000e+00,  0.00000000e+00, -0.00000000e+00,
                 8.80751720e-01, -9.97142430e-01],
               [ 9.99948603e-01, -9.99948603e-01,  1.19196784e-01,
                -2.47166897e-03,  8.80751720e-01, -9.97142430e-01,
                 9.91155487e-08, -3.34504020e-04]]])

B = np.array([[[0., 1.],
               [2., 1.]],

              [[0., 1.],
               [1., 0.]]])

C = np.empty((A.shape[-1], B.shape[-1])

for i in range(C.shape[0]):
    for j in range(C.shape[1]):
        C[i, j] = (A[..., i] * B[..., j]).sum()