np.apply_along_axis
на самом деле вам не поможет, потому что вы пытаетесь перебрать два массива.По сути, вам придется использовать цикл, как описано здесь .
Теперь, циклы в порядке, если ваши массивы маленькие, но если N и P большие, то вы, вероятно, захотитевместо этого использовать FFT для свертки.
Однако, вам нужно сначала обнулить массивы, чтобы ваша «полная» свертка имела ожидаемую форму:
M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)
A_ = np.zeros((M, N+P-1), dtype=A.dtype)
A_[:, :N] = A
B_ = np.zeros((M, N+P-1), dtype=B.dtype)
B_[:, :P] = B
A_fft = np.fft.fft(A_, axis=1)
B_fft = np.fft.fft(B_, axis=1)
C_fft = A_fft * B_fft
C = np.real(np.fft.ifft(C_fft))
# Test
C_test = np.zeros((M, N+P-1))
for i in range(M):
C_test[i, :] = np.convolve(A[i, :], B[i, :], 'full')
assert np.allclose(C, C_test)