How do I do an einsum that mimics 'keepdims'?

Question

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?

Answer 1

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

Adding a dimension in numpy is at worst O(ndim), so basically free. Crucially, the actually data is not touched - all that happens is that the .strides and .shape tuples get an extra element each

There is no way right now to use einsum to directly get what you want.

You could try to make a pull-request against numpy to support something like ijk, ijk -> i1k, if you really think it improves readability