Search code examples
pythonarraysnumpylinear-algebrainner-product

Compute numpy.inner() over first (instead of last) axis

I'm trying to make a function like numpy.inner, but which sums over the first axis of both arrays instead of the last axis. Currently I'm using tensordot with rollaxis:

def inner1(a, b):
    return numpy.tensordot(numpy.rollaxis(a, 0, len(a.shape)), b, 1)

but I'm wondering: is there a better way? Perhaps one that doesn't require me to roll the axes?

I feel like einsum should make this possible, but I'm not sure how to use it here.
It seems to require me to hard-code the dimensionality of a and b when I specify the subscripts string, which I can't really do here because there is no particular requirement on the input dimensionality.

(Note: I am aware that there are performance implications to summing over the first axis instead of the last, but I'm ignoring them here.)

Solution

  • I think what you want is np.tensordot(a, b, (0, 0)).