code: def _fft_lowering_cpu(ctx, x, *, fft_type, fft_lengths):
      (x_aval,) = ctx.avals_in
  
      in_shape = x_aval.shape
      dtype = x_aval.dtype
      (out_aval,) = ctx.avals_out
      out_shape = out_aval.shape
  
      forward = fft_type in (
          xla_client.FftType.FFT,
          xla_client.FftType.RFFT,
      )
      ndims = len(in_shape)
      assert len(fft_lengths) >= 1
      assert len(fft_lengths) <= ndims, (fft_lengths, ndims)
      assert len(in_shape) == len(out_shape) == ndims
  
      # PocketFft does not allow size 0 dimensions.
      if 0 in in_shape or 0 in out_shape:
          if fft_type == xla_client.FftType.RFFT:
              assert dtype in (np.float32, np.float64), dtype
              out_dtype = np.dtype(
                  np.complex64 if dtype == np.float32 else np.complex128
              )
  
          elif fft_type == xla_client.FftType.IRFFT:
              assert np.issubdtype(dtype, np.complexfloating), dtype
              out_dtype = np.dtype(
                  np.float32 if dtype == np.complex64 else np.float64
              )
  
          else:
              assert np.issubdtype(dtype, np.complexfloating), dtype
              out_dtype = dtype
  
          zero = mlir.ir_constant(np.array(0, dtype=out_dtype))
          return [
              mlir.broadcast_in_dim(
                  ctx, zero, out_aval, broadcast_dimensions=[]
              )
          ]
  
      strides_in = []
      stride = 1
      for d in reversed(in_shape):
          strides_in.append(stride)
          stride *= d
      strides_in = mlir.shape_tensor(
          mlir.eval_dynamic_shape(ctx, tuple(reversed(strides_in)))
      )
  
      strides_out = []
      stride = 1
      for d in reversed(out_shape):
          strides_out.append(stride)
          stride *= d
      strides_out = mlir.shape_tensor(
          mlir.eval_dynamic_shape(ctx, tuple(reversed(strides_out)))
      )
  
      # scale = 1. if forward else (1. / np.prod(fft_lengths)) as a f64[1] tensor
      double_type = mlir.ir.RankedTensorType.get(
          (), mlir.ir.F64Type.get()
      )
      size_fft_length_prod = np.prod(fft_lengths) if fft_lengths else 1
      (size_fft_lengths,) = mlir.eval_dynamic_shape_as_vals(
          ctx, (size_fft_length_prod,)
      )
      size_fft_lengths = hlo.ConvertOp(double_type, size_fft_lengths)
      one = mlir.ir_constant(np.float64(1.0))
      scale = one if forward else hlo.DivOp(one, size_fft_lengths)
      scale = hlo.ReshapeOp(
          mlir.ir.RankedTensorType.get((1,), mlir.ir.F64Type.get()),
          scale,
      ).result
  
      in_shape = mlir.shape_tensor(
          mlir.eval_dynamic_shape(ctx, in_shape)
      )
      out_shape = mlir.shape_tensor(
          mlir.eval_dynamic_shape(ctx, out_shape)
      )
      in_shape = (
          in_shape
          if fft_type != xla_client.FftType.IRFFT
          else out_shape
      )
  
      result_type = mlir.aval_to_ir_type(out_aval)
      return [
          ducc_fft.dynamic_ducc_fft_hlo(
              result_type,
              x,
              input_dtype=x_aval.dtype,
              ndims=ndims,
              input_shape=in_shape,
              strides_in=strides_in,
              strides_out=strides_out,
              scale=scale,
              fft_type=fft_type,
              fft_lengths=fft_lengths,
              result_shape=out_shape,
          )
      ]
  
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 6 for calling a global function with example usage: _fft_lowering_cpu(ctx, x) and returns: list, list