code: def _dct_ortho_norm(out: Array, axis: int) -> Array:
      factor = lax.concatenate(
          [
              lax.full((1,), 4, out.dtype),
              lax.full((out.shape[axis] - 1,), 2, out.dtype),
          ],
          0,
      )
      factor = lax.expand_dims(
          factor, [a for a in range(out.ndim) if a != axis]
      )
      return out / lax.sqrt(factor * out.shape[axis])
  
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 3 for calling a global function with example usage: _dct_ortho_norm(out, axis) and returns: unknown