code: def _expand_to_beam_size(
      data, beam_size, batch_size, state_batch_axis=None
  ):
      """Tile all the states to have batch_size * beam_size on the batch axis.
  
      Parameters
      ----------
      data : A single mx.np.ndarray or nested container with mx.np.ndarray
          Each mx.np.ndarray should have shape (N, ...) when state_info is None,
          or same as the layout in state_info when it's not None.
      beam_size : int
          Beam size
      batch_size : int
          Batch size
      state_batch_axis : Nested structure of dictionary, default None.
          Descriptors for states, usually from decoder's ``state_batch_axis()``.
          When None, this method assumes that the batch axis is the first dimension.
      Returns
      -------
      new_states : Object that contains mx.np.ndarray
          Each mx.np.ndarray should have shape batch_size * beam_size on the batch axis.
      """
      if isinstance(data, (list, tuple)):
          if state_batch_axis is not None:
              # TODO(sxjscience) Better Exception Handling
              return [
                  _expand_to_beam_size(
                      d, beam_size, batch_size, batch_axis
                  )
                  for d, batch_axis in zip(data, state_batch_axis)
              ]
          else:
              return [
                  _expand_to_beam_size(d, beam_size, batch_size, None)
                  for d in data
              ]
      elif isinstance(data, dict):
          if state_batch_axis is not None:
              return {
                  k: _expand_to_beam_size(
                      v, beam_size, batch_size, state_batch_axis[k]
                  )
                  for k, v in data.items()
              }
          else:
              return {
                  k: _expand_to_beam_size(
                      v, beam_size, batch_size, None
                  )
                  for k, v in data.items()
              }
      elif isinstance(data, mx.np.ndarray):
          if state_batch_axis is None:
              batch_axis = 0
          else:
              batch_axis = state_batch_axis
          if data.shape[batch_axis] != batch_size:
              raise ValueError(
                  "The batch size of all the inner elements in states must be "
                  "{}, Found shape={}, inferred batch axis={}".format(
                      batch_size, data.shape, batch_axis
                  )
              )
          new_shape = list(data.shape)
          new_shape[batch_axis] = batch_size * beam_size
          new_shape = tuple(new_shape)
          bcast_new_shape = (
              new_shape[:batch_axis]
              + (batch_size, beam_size)
              + new_shape[(batch_axis + 1) :]
          )
          return (
              mx.np.expand_dims(data, batch_axis + 1)
              .broadcast_to(bcast_new_shape)
              .reshape(new_shape)
          )
      elif data is None:
          return None
      else:
          raise NotImplementedError
  
imports: import warnings
  import mxnet
  import abc
  from mxnet.gluon import HybridBlock
  from typing import Callable, Optional
  from .layers import get_activation
label_desc: Use image node_id 1 for calling a global function with example usage: _expand_to_beam_size(data, beam_size, batch_size, state_batch_axis) and returns: unknown, unknown, unknown, unknown, mx, None