4

a python question: I've got a np.einsum operation that I'm doing on a pair of 3d arrays:

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

Problem I'm having is the result is 2d; the operation collapses the 'j' dimension. What I'd love to do is to have it retain the 'j' dimension similar to how 'keepdims' works in the np.sum function.

I can wrap the result in np.expand_dims, but that seems inefficient to me. I'd prefer to find some way to tweak the einsum to output what I'm after.

Is this posible?

4

1 回答 1

7

我可以将结果包装在 中np.expand_dims，但这对我来说似乎效率低下

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

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

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

于 2017-03-23T18:09:42.300 回答