code: def fftshift(
      x: ArrayLike, axes: None | int | Sequence[int] = None
  ) -> Array:
      check_arraylike("fftshift", x)
      x = jnp.asarray(x)
      shift: int | Sequence[int]
      if axes is None:
          axes = tuple(range(x.ndim))
          shift = [dim // 2 for dim in x.shape]
      elif isinstance(axes, int):
          shift = x.shape[axes] // 2
      else:
          shift = [x.shape[ax] // 2 for ax in axes]
  
      return jnp.roll(x, shift, axes)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import operator
  import numpy
  from jax import dtypes
  from jax import lax
  from jax._src.lib import xla_client
  from jax._src.util import safe_zip
  from jax._src.numpy.util import check_arraylike, _wraps
  from jax._src.numpy import lax_numpy
  from jax._src.numpy import ufuncs, reductions
  from jax._src.typing import Array, ArrayLike
label_desc: Use image node_id 22 for calling a global function with example usage: fftshift(x, axes) and returns: jnp