code: def maybe_get_axis(axis: int, arrays: Any) -> Optional[int]:
      """Returns `array.shape[axis]` for one of the arrays in the input."""
      shapes = [a.shape for a in jax.tree_util.tree_leaves(arrays)]
      sizes = {s[axis] for s in shapes}
      if len(sizes) != 1:
          raise ValueError(
              "Arrays must have the same mapped axis size, found "
              f"sizes {sizes} for input shapes {shapes}"
          )
      (size,) = sizes
      return size
  
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 22 for calling a global function with example usage: maybe_get_axis(axis, arrays) and returns: size