jax.scipy.fft.dctn#

jax.scipy.fft.dctn(x, type=2, s=None, axes=None, norm=None)[source]#

Computes the multidimensional discrete cosine transform of the input

JAX implementation of scipy.fft.dctn().

Parameters:

  • x (Array) – array

  • type (int) – integer, default = 2. Currently only type 2 is supported.

  • s (Sequence[int] | None) – integer or sequence of integers. Specifies the shape of the result. If not specified, it will default to the shape of x along the specified axes.

  • axes (Sequence[int] | None) – integer or sequence of integers. Specifies the axes along which the transform will be computed.

  • norm (str | None) – string. The normalization mode: one of [None, "backward", "ortho"]. The default is None, which is equivalent to "backward".

Returns:

array containing the discrete cosine transform of x

Return type:

Array

See also

Examples

jax.scipy.fft.dctn computes the transform along both the axes by default when axes argument is None.

>>> x = jax.random.normal(jax.random.key(0), (3, 3))
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dctn(x))
[[ 12.01   6.2  -10.17]
 [  8.84   9.65  -3.54]
 [ 11.25  -1.54  -0.88]]

When s=[2], dimension of the transform along axis 0 will be 2 and dimension along axis 1 will be same as that of input.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dctn(x, s=[2]))
[[ 9.36 10.22 -8.53]
 [11.57  2.85 -2.06]]

When s=[2] and axes=[1], dimension of the transform along axis 1 will be 2 and dimension along axis 0 will be same as that of input. Also when axes=[1], transform will be computed only along axis 1.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dctn(x, s=[2], axes=[1]))
[[ 7.3  -0.57]
 [ 0.19 -0.36]
 [-0.   -1.4 ]]

When s=[2, 4], shape of the transform will be (2, 4).

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dctn(x, s=[2, 4]))
[[  9.36  11.23   2.12 -10.97]
 [ 11.57   5.86  -1.37  -1.58]]