Вычислить numpy.inner() по первой (а не по последней) оси

Question

Я пытаюсь сделать такую ​​функцию numpy.inner, но который суммируется по первой оси обоих массивов вместо последней оси. В настоящее время я использую tensordot с rollaxis:

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

но мне интересно: есть ли лучший способ? Возможно тот, который не требует, чтобы я катил топоры?

я чувствую einsum должно сделать это возможным, но я не уверен, как использовать это здесь.
Кажется, мне нужно жестко закодировать размерность a а также b когда я указываю строку подписки, что я не могу здесь сделать, потому что нет особых требований к входной размерности.

(Примечание: я знаю, что производительность суммируется по первой оси вместо последней, но здесь я их игнорирую.)

Answer 1

Я думаю, что вы хотите np.tensordot(a, b, (0, 0)),

Answer 2

Это не так красиво, как tensordot решение, но вы можете построить einsum строка из ndim из входов:

ll = 'abcdefghijklmnopqrstuvw'
astr = ll[0]+ll[1:a.ndim]+','+ll[0]+ll[a.ndim:a.ndim+b.ndim-1]
np.einsum(astr,a,b)

np.einsum позволяет указывать оси в виде списков, а не строки

np.einsum(a, [0]+range(1,a.ndim), b, [0]+range(a.ndim,a.ndim+b.ndim-1))

Для пары массивов 3d и 2d они производят:

 np.einsum('abc,ad', a, b)
 np.einsum(a, [0,1,2], b, [0,3])

'...' здесь не работает, потому что это подразумевает повторяющиеся оси (насколько это возможно), где вы хотите уникальные оси (кроме 1-й).

Хотя писать сложнее, einsum решение быстрее, чем tensordot один (в 3 раза быстрее для небольших тестовых массивов).

---

Еще один вариант с einsum это изменить форму массивов, уменьшив "оставшиеся" размеры до одного. Это добавляет немного времени к расчету, но не намного:

np.einsum('ij,ik',a.reshape(a.shape[0],-1), b.reshape(a.shape[0],-1)).reshape(a.shape[1:]+b.shape[1:])