python - 我如何做一个模仿 'keepdims' 的 einsum？

Question

一个 python 问题:我有一个 np.einsum 操作，我正在一对 3d 数组上做:

return np.einsum('ijk, ijk -> ik', input_array, self._beta_array)

我遇到的问题是结果是 2d；该操作折叠了“j”维度。我想做的是让它保留“j”维度，类似于“keepdims”在 np.sum 函数中的工作方式。

我可以将结果包装在 np.expand_dims 中，但这对我来说似乎效率低下。我更愿意找到一些方法来调整 einsum 以输出我想要的内容。

这可能吗？

Answer 1

I can wrap the result in np.expand_dims, but that seems inefficient to me

在 numpy 中添加一个维度最差 O(ndim)，所以基本上是免费的。至关重要的是，实际的数据并没有被触及——所发生的只是 .strides 和 .shape 元组各自获得一个额外的元素

目前没有办法使用einsum直接得到你想要的。

如果你真的认为它提高了可读性，你可以尝试对 numpy 发出拉取请求以支持像 ijk, ijk -> i1k 这样的东西