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.)
I think what you want is np.tensordot(a, b, (0, 0)).