code: def dctn(
      x: Array,
      type: int = 2,
      s: Sequence[int] | None = None,
      axes: Sequence[int] | None = None,
      norm: str | None = None,
  ) -> Array:
      if type != 2:
          raise NotImplementedError("Only DCT type 2 is implemented.")
  
      if axes is None:
          axes = range(x.ndim)
  
      if len(axes) == 1:
          return dct(
              x,
              n=s[0] if s is not None else None,
              axis=axes[0],
              norm=norm,
          )
  
      if s is not None:
          ns = dict(zip(axes, s))
          pads = [
              (0, ns[a] - x.shape[a] if a in ns else 0, 0)
              for a in range(x.ndim)
          ]
          x = lax.pad(x, jnp.array(0, x.dtype), pads)
  
      if len(axes) == 2:
          return _dct2(x, axes=axes, norm=norm)
  
      # compose high-D DCTs from 2D and 1D DCTs:
      for axes_block in [
          axes[i : i + 2] for i in range(0, len(axes), 2)
      ]:
          x = dctn(x, axes=axes_block, norm=norm)
      return x
  
file_path: jax/jax/_src/scipy/fft.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import math
  import scipy.fft
  from jax import lax
  import jax.numpy
  from jax._src.util import canonicalize_axis
  from jax._src.numpy.util import _wraps, promote_dtypes_complex
  from jax._src.typing import Array
label_desc: Use image node_id 6 for calling a global function with example usage: dctn(x, type, s, axes, norm) and returns: x, dct, _dct2