jax.scipy.fft.dct#

jax.scipy.fft.dct(x, type=2, n=None, axis=-1, norm=None)[source]#

Computes the discrete cosine transform of the input

JAX implementation of scipy.fft.dct().

Parameters:
  • x (Array) – array

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

  • n (int | None | None) – integer, default = x.shape[axis]. The length of the transform. If larger than x.shape[axis], the input will be zero-padded, if smaller, the input will be truncated.

  • axis (int) – integer, default=-1. The axis along which the dct will be performed.

  • norm (str | None | 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

>>> x = jax.random.normal(jax.random.key(0), (3, 3))
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dct(x))
[[ 6.43  3.56 -2.86]
 [-1.75  1.55 -1.4 ]
 [ 1.33 -2.01 -0.82]]

When n smaller than x.shape[axis]

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

When n smaller than x.shape[axis] and axis=0

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dct(x, n=2, axis=0))
[[ 3.09  4.4  -2.81]
 [ 2.41  2.62  0.76]]

When n larger than x.shape[axis] and axis=1

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jax.scipy.fft.dct(x, n=4, axis=1))
[[ 6.43  4.88  0.04 -3.3 ]
 [-1.75  0.73  1.01 -2.18]
 [ 1.33 -1.05 -2.34 -0.07]]