jax.scipy.signal.fftconvolve

jax.scipy.signal.fftconvolve(in1, in2, mode='full', axes=None)

使用快速傅里叶变换 (FFT) 对两个 N 维数组进行卷积。

JAX 对 scipy.signal.fftconvolve() 的实现。

参数:

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

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

• mode (str) –

  控制输出的大小。可用操作包括

  • "full": （默认）输出输入的完整卷积。

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

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

• axes (Sequence[int] | None) – 可选的轴序列，沿这些轴应用卷积。

返回:

包含卷积结果的数组。

返回类型:

Array

另请参阅

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

• jax.scipy.signal.convolve(): 直接卷积

示例

一些一维卷积的示例。由于基于 FFT 的卷积是近似的，因此我们在下面使用 jax.numpy.printoptions() 来调整打印精度

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

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

>>> with jax.numpy.printoptions(precision=3):
...   print(jax.scipy.signal.fftconvolve(x, y, mode='full'))
[1. 3. 6. 7. 6. 3. 1.]

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

>>> with jax.numpy.printoptions(precision=3):
...   print(jax.scipy.signal.fftconvolve(x, y, mode='same'))
[3. 6. 7. 6. 3.]

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

>>> with jax.numpy.printoptions(precision=3):
...   print(jax.scipy.signal.fftconvolve(x, y, mode='valid'))
[6. 7. 6.]