向量化autograd按行卷积

Question

我想使用autograd来确定涉及卷积的函数的梯度。我想要的是类似的东西

# preparation
import numpy as np
import autograd.numpy as ag
arr1 = np.random.rand(4000, 50)
arr2 = np.random.rand(4000, 50)

# the operation of interest
result = [ag.convolve(a, b) for a, b in zip(arr1, arr2)]

上述解决方案有效，但效率低下且不明确。我想对操作进行矢量化处理，因此只需要一个函数调用autograd。在纯numpy / scipy中，可以通过指定轴参数轻松完成此操作：

from scipy import signal

result = signal.oaconvolve(arr1, arr2, axes=1)

Autograd为convolve方法提供了一个axes参数。经过大量的反复试验并阅读了this问题，我发现

from autograd.scipy import signal as agsignal

result = agsignal.convolve(arr1, arr2, axes=([1], [1]))

几乎可以满足我的需求。但是，以上命令返回arr1的每一行与arr2的每一行进行卷积，而我想对{{1}的第n行进行卷积}}和arr1的第n行。我该怎么办？