code: def has_batch_size(
      batch_elems: tp.Tuple[tp.Any, tp.Any, tp.Any], batch_size: int
  ) -> bool:
      sizes: tp.Set[int] = set(
          x.shape[0] for x in jax.tree_leaves(batch_elems)
      )
  
      if len(sizes) == 0:
          return True
      elif len(sizes) > 1:
          raise ValueError(
              f"Data cardinality is ambiguous, got multiple batch dimension: {sizes}"
          )
  
      current_batch_size = sizes.pop()
  
      return current_batch_size == batch_size
  
imports: import collections
  import functools
  import math
  import typing
  import jax
  import jax.numpy
  import numpy
  import six
  import typing_extensions
  from elegy import types, utils
label_desc: Use image node_id 11 for calling a global function with example usage: has_batch_size(batch_elems, batch_size) and returns: unknown, True