code: def to_axes_or_slice(axis: AxisOrAxes) -> AxesOrSlice:
      if isinstance(axis, slice):
          return axis
      elif isinstance(axis, int):
          return (axis,)
      elif isinstance(axis, collections.abc.Iterable) and all(
          isinstance(ax, int) for ax in axis
      ):
          return tuple(axis)
      else:
          raise ValueError(
              f"`axis` should be an int, slice or iterable of ints. Got: {axis}"
          )
  
imports: import collections.abc
  from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: to_axes_or_slice(axis) and returns: axis, axis, tuple