code: def fft(
      x, fft_type: xla_client.FftType | str, fft_lengths: Sequence[int]
  ):
      if isinstance(fft_type, str):
          typ = _str_to_fft_type(fft_type)
      elif isinstance(fft_type, xla_client.FftType):
          typ = fft_type
      else:
          raise TypeError(f"Unknown FFT type value '{fft_type}'")
  
      if typ == xla_client.FftType.RFFT:
          if np.iscomplexobj(x):
              raise ValueError(
                  "only real valued inputs supported for rfft"
              )
          (x,) = promote_dtypes_inexact(x)
      else:
          (x,) = promote_dtypes_complex(x)
      if len(fft_lengths) == 0:
          # XLA FFT doesn't support 0-rank.
          return x
      fft_lengths = tuple(fft_lengths)
      return fft_p.bind(x, fft_type=typ, fft_lengths=fft_lengths)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import math
  import numpy
  from jax import lax
  from jax._src import dispatch
  from jax._src.api import jit, linear_transpose, ShapeDtypeStruct
  from jax._src.core import Primitive, is_constant_shape
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib import xla_client
  from jax._src.lib import xla_extension_version
  from jax._src.lib import ducc_fft
  from jax._src.numpy.util import promote_dtypes_complex, promote_dtypes_inexact
label_desc: Use image node_id 2 for calling a global function with example usage: fft(x, fft_type, fft_lengths) and returns: fft_p, x