def _expand_to_beam_size(data, beam_size, batch_size, state_info=None):
assert not state_info or isinstance(state_info, (type(data), dict)),'<STR_LIT>''<STR_LIT>'.format(type(state_info), type(data))<EOL>if isinstance(data, list):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = [None] * len(data)<EOL><DEDENT>return [_expand_to_beam_size(d, beam_size, batch_size, s)<EOL>for d, s in zip(data, state_info)]<EOL><DEDENT>elif isinstance(data, tuple):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = [None] * len(data)<EOL>state_info = tuple(state_info)<EOL><DEDENT>return tuple(_expand_to_beam_size(d, beam_size, batch_size, s)<EOL>for d, s in zip(data, state_info))<EOL><DEDENT>elif isinstance(data, dict):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = {k: None for k in data.keys()}<EOL><DEDENT>return {k: _expand_to_beam_size(v, beam_size, batch_size, state_info[k])<EOL>for k, v in data.items()}<EOL><DEDENT>elif isinstance(data, mx.nd.NDArray):<EOL><INDENT>if not state_info:<EOL><INDENT>batch_axis = <NUM_LIT:0><EOL><DEDENT>else:<EOL><INDENT>batch_axis = state_info['<STR_LIT>'].find('<STR_LIT:N>')<EOL><DEDENT>if data.shape[batch_axis] != batch_size:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(batch_size, data.shape))<EOL><DEDENT>new_shape = list(data.shape)<EOL>new_shape[batch_axis] = batch_size * beam_size<EOL>new_shape = tuple(new_shape)<EOL>return data.expand_dims(batch_axis+<NUM_LIT:1>).broadcast_axes(axis=batch_axis+<NUM_LIT:1>, size=beam_size).reshape(new_shape)<EOL><DEDENT>elif isinstance(data, mx.sym.Symbol):<EOL><INDENT>if not state_info:<EOL><INDENT>batch_axis = <NUM_LIT:0><EOL><DEDENT>else:<EOL><INDENT>batch_axis = state_info['<STR_LIT>'].find('<STR_LIT:N>')<EOL><DEDENT>new_shape = (<NUM_LIT:0>, ) * batch_axis + (-<NUM_LIT:3>, -<NUM_LIT:2>)<EOL>return data.expand_dims(batch_axis+<NUM_LIT:1>).broadcast_axes(axis=batch_axis+<NUM_LIT:1>, size=beam_size).reshape(new_shape)<EOL><DEDENT>else:<EOL><INDENT>raise NotImplementedError<EOL><DEDENT>
Tile all the states to have batch_size * beam_size on the batch axis.

    Parameters
    ----------
    data : A single NDArray/Symbol or nested container with NDArrays/Symbol
        Each NDArray/Symbol 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_info : Nested structure of dictionary, default None.
        Descriptors for states, usually from decoder's ``state_info()``.
        When None, this method assumes that the batch axis is the first dimension.
    Returns
    -------
    new_states : Object that contains NDArrays/Symbols
        Each NDArray/Symbol should have shape batch_size * beam_size on the batch axis.