code: def _get_fftn_test_axes(shape):
      axes = [[]]
      ndims = len(shape)
      # XLA's FFT op only supports up to 3 innermost dimensions.
      if ndims <= 3:
          axes.append(None)
      for naxes in range(1, min(ndims, 3) + 1):
          axes.extend(itertools.combinations(range(ndims), naxes))
      for index in range(1, ndims + 1):
          axes.append((-index,))
      return axes
  
imports: import itertools
  import numpy
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import lax
  from jax import numpy
  from jax._src import config
  from jax._src import dtypes
  from jax._src import test_util
  from jax._src.numpy.util import promote_dtypes_complex
label_desc: Use image node_id 1 for calling a global function with example usage: _get_fftn_test_axes(shape) and returns: axes