numpy einsum:嵌套点积 [英] numpy einsum: nested dot products

问题描述

我有两个n -by- k -by- 3数组a和b，例如

I have two n-by-k-by-3 arrays a and b, e.g.,

import numpy as np

a = np.array([
    [
        [1, 2, 3],
        [3, 4, 5]
    ],
    [
        [4, 2, 4],
        [1, 4, 5]
    ]
    ])
b = np.array([
    [
        [3, 1, 5],
        [0, 2, 3]
    ],
    [
        [2, 4, 5],
        [1, 2, 4]
    ]
    ])

，并且它想计算所有三胞胎"对的点积，即

and it like to compute the dot-product of all pairs of "triplets", i.e.,

np.sum(a*b, axis=2)

一种更好的方法可能是 einsum ，但我似乎无法弄清楚索引.

A better way to do that is perhaps einsum, but I can't seem to get the indices straight.

这里有任何提示吗?

推荐答案

您正在通过减总和来减少这两个3D输入数组上的第三个轴，同时保持前两个轴对齐.因此，使用 np.einsum ，我们可以前两个字符串相同，第三个字符串也相同，但是在输出字符串表示法中将被跳过，我们将沿两个输入的该轴递减.因此，解决方案将是-

You are loosing the third axis on those two 3D input arrays with that sum-reduction, while keeping the first two axes aligned. Thus, with np.einsum, we would have the first two strings identical alongwith the third string being identical too, but would be skipped in the output string notation signalling we are reducing along that axis for both the inputs. Thus, the solution would be -

np.einsum('ijk,ijk->ij',a,b)

这篇关于numpy einsum:嵌套点积的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！