code: def get_mapped_axis_size(args: tuple[Any], in_axes: Any) -> int:
      sizes = uniq(
          jax.tree_util.tree_leaves(
              jax.tree_util.tree_map(maybe_get_axis, in_axes, args)
          )
      )
      assert sizes, "hk.vmap should guarantee non-empty in_axes"
      # NOTE: We use the first in_axes regardless of how many non-unique values
      # there are to allow JAX to handle multiple conflicting sizes.
      return sizes[0]
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 23 for calling a global function with example usage: get_mapped_axis_size(args, in_axes) and returns: sizes