jax.numpy.fft.fft2

jax.numpy.fft.fft2(a, s=None, axes=(-2, -1), norm=None)

Compute a two-dimensional discrete Fourier transform along given axes.

JAX implementation of numpy.fft.fft2().

Parameters:

• a (ArrayLike) – input array. Must have a.ndim >= 2.

• s (Shape | None | None) – optional length-2 sequence of integers. Specifies the size of the output along each specified axis. If not specified, it will default to the size of a along the specified axes.

• axes (Sequence[int]) – optional length-2 sequence of integers, default=(-2,-1). Specifies the axes along which the transform is computed.

• norm (str | None | None) – string, default=”backward”. The normalization mode. “backward”, “ortho” and “forward” are supported.

Returns:

An array containing the two-dimensional discrete Fourier transform of a along given axes.

Return type:

Array

See also

• jax.numpy.fft.fft(): Computes a one-dimensional discrete Fourier transform.

• jax.numpy.fft.fftn(): Computes a multidimensional discrete Fourier transform.

• jax.numpy.fft.ifft2(): Computes a two-dimensional inverse discrete Fourier transform.

Examples

jnp.fft.fft2 computes the transform along the last two axes by default.

>>> x = jnp.array([[[1, 3],
...                 [2, 4]],
...                [[5, 7],
...                 [6, 8]]])
>>> with jnp.printoptions(precision=2, suppress=True):
...   jnp.fft.fft2(x)
Array([[[10.+0.j, -4.+0.j],
        [-2.+0.j,  0.+0.j]],

       [[26.+0.j, -4.+0.j],
        [-2.+0.j,  0.+0.j]]], dtype=complex64)

When s=[2, 3], dimension of the transform along axes (-2, -1) will be (2, 3) and dimension along other axes will be the same as that of input.

>>> with jnp.printoptions(precision=2, suppress=True):
...   jnp.fft.fft2(x, s=[2, 3])
Array([[[10.  +0.j  , -0.5 -6.06j, -0.5 +6.06j],
        [-2.  +0.j  , -0.5 +0.87j, -0.5 -0.87j]],

       [[26.  +0.j  ,  3.5-12.99j,  3.5+12.99j],
        [-2.  +0.j  , -0.5 +0.87j, -0.5 -0.87j]]], dtype=complex64)

When s=[2, 3] and axes=(0, 1), shape of the transform along axes (0, 1) will be (2, 3) and dimension along other axes will be same as that of input.

>>> with jnp.printoptions(precision=2, suppress=True):
...   jnp.fft.fft2(x, s=[2, 3], axes=(0, 1))
Array([[[14. +0.j  , 22. +0.j  ],
        [ 2. -6.93j,  4.-10.39j],
        [ 2. +6.93j,  4.+10.39j]],

       [[-8. +0.j  , -8. +0.j  ],
        [-2. +3.46j, -2. +3.46j],
        [-2. -3.46j, -2. -3.46j]]], dtype=complex64)

jnp.fft.ifft2 can be used to reconstruct x from the result of jnp.fft.fft2.

>>> x_fft2 = jnp.fft.fft2(x)
>>> jnp.allclose(x, jnp.fft.ifft2(x_fft2))
Array(True, dtype=bool)