code: def _dct2(x: Array, axes: Sequence[int], norm: str | None) -> Array:
      axis1, axis2 = map(
          partial(canonicalize_axis, num_dims=x.ndim), axes
      )
      N1, N2 = x.shape[axis1], x.shape[axis2]
      v = _dct_interleave(_dct_interleave(x, axis1), axis2)
      V = jnp.fft.fftn(v, axes=axes)
      k1 = lax.expand_dims(
          jnp.arange(N1, dtype=V.dtype),
          [a for a in range(x.ndim) if a != axis1],
      )
      k2 = lax.expand_dims(
          jnp.arange(N2, dtype=V.dtype),
          [a for a in range(x.ndim) if a != axis2],
      )
      out = _W4(N1, k1) * (
          _W4(N2, k2) * V
          + _W4(N2, -k2)
          * jnp.roll(jnp.flip(V, axis=axis2), shift=1, axis=axis2)
      )
      out = 2 * out.real
      if norm == "ortho":
          return _dct_ortho_norm(_dct_ortho_norm(out, axis1), axis2)
      return out
  
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 5 for calling a global function with example usage: _dct2(x, axes, norm) and returns: out, _dct_ortho_norm