code: def is_match_states_batch_size(
      states, states_batch_axis, batch_size
  ) -> bool:
      """Test whether the generated states have the specified batch size
  
      Parameters
      ----------
      states
          The states structure
      states_batch_axis
          The states batch axis structure
      batch_size
          The batch size
  
      Returns
      -------
      ret
      """
      if states_batch_axis is None:
          return True
      if isinstance(states_batch_axis, int):
          if states.shape[states_batch_axis] == batch_size:
              return True
      for ele_states_batch_axis, ele_states in zip(
          states_batch_axis, states
      ):
          ret = is_match_states_batch_size(
              ele_states, ele_states_batch_axis, batch_size
          )
          if ret is False:
              return False
      return True
  
file_path: gluon-nlp/src/gluonnlp/utils/testing.py
imports: import numpy.testing
  import numpy
  import mxnet
  from mxnet.util import use_np
  from .parameter import move_to_ctx
label_desc: Use image node_id 1 for calling a global function with example usage: is_match_states_batch_size(states, states_batch_axis, batch_size) and returns: True, True, True, False