jax.scipy.signal.convolve

jax.scipy.signal.convolve(in1, in2, mode='full', method='auto', precision=None)

两个 N 维数组的卷积。

scipy.signal.convolve() 的 JAX 实现。

参数:

• in1 (Array) – 卷积的左侧输入。

• in2 (Array) – 卷积的右侧输入。必须满足 in1.ndim == in2.ndim。

• mode (str) –

  控制输出的大小。可用的操作有：

  • "full"： (默认) 输出输入的完整卷积。

  • "same"：返回 "full" 输出的中心部分，该部分的大小与 in1 相同。

  • "valid"：返回 "full" 输出中不依赖于数组边缘填充的部分。

• method (str) –

  控制计算方法。选项有：

  • "auto"： (默认) 始终使用 "direct" 方法。

  • "direct"：降级到 jax.lax.conv_general_dilated()。

  • "fft"：通过快速傅里叶变换计算结果。

• precision (PrecisionLike | None) – 指定计算的精度。有关可用值的描述，请参阅 jax.lax.Precision。

返回值:

包含卷积结果的数组。

返回类型:

Array

另请参阅

• jax.numpy.convolve()：一维卷积

• jax.scipy.signal.convolve2d()：二维卷积

• jax.scipy.signal.correlate()：多维相关

示例

一些一维卷积示例

>>> x = jnp.array([1, 2, 3, 2, 1])
>>> y = jnp.array([1, 1, 1])

完整卷积在边缘使用隐式零填充

>>> jax.scipy.signal.convolve(x, y, mode='full')
Array([1., 3., 6., 7., 6., 3., 1.], dtype=float32)

指定 mode = 'same' 返回与第一个输入大小相同的中心卷积

>>> jax.scipy.signal.convolve(x, y, mode='same')
Array([3., 6., 7., 6., 3.], dtype=float32)

指定 mode = 'valid' 仅返回两个数组完全重叠的部分

>>> jax.scipy.signal.convolve(x, y, mode='valid')
Array([6., 7., 6.], dtype=float32)