code: def _resolve_axis(
      axis_spec: AxisSpec, elem: Any
  ) -> PyTree[ResolvedAxisSpec]:
      if axis_spec is None or isinstance(axis_spec, int):
          return jtu.tree_map(lambda _: axis_spec, elem)
      elif callable(axis_spec):
          return jtu.tree_map(axis_spec, elem)
      else:
          raise ValueError(
              "`in_axes` and `out_axes` must consist of None, ints, and callables only."
          )
  
imports: import dataclasses
  import functools
  import inspect
  import warnings
  from collections.abc import Callable, Hashable
  from typing import Any, Literal, Optional, overload, Union
  import jax
  import jax._src.traceback_util
  import jax.core
  import jax.interpreters.batching
  import jax.interpreters.pxla
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import PyTree
  from ._compile_utils import compile_cache, hashable_combine, hashable_partition, Lowered
  from ._custom_types import sentinel
  from ._deprecate import deprecated_0_10
  from ._doc_utils import doc_remove_args
  from ._filters import combine, filter, is_array, is_array_like, partition
  from ._module import Module, module_update_wrapper, Partial, Static
label_desc: Use image node_id 2 for calling a global function with example usage: _resolve_axis(axis_spec, elem) and returns: jtu, jtu