import torch


def create_lm_batch_state(lm_states_list, lm_layers, is_wordlm):
    """Create batch of LM states.

    Args:
        lm_states (list or dict): list of individual LM states
        lm_layers (int): number of LM layers
        is_wordlm (bool): whether provided LM is a word-LM

    Returns:
       batch_states (list): batch of LM states

    """
    if is_wordlm:
        batch_states = lm_states_list
    else:
        batch_states = {}

        batch_states["c"] = [
            torch.stack([state["c"][layer] for state in lm_states_list])
            for layer in range(lm_layers)
        ]
        batch_states["h"] = [
            torch.stack([state["h"][layer] for state in lm_states_list])
            for layer in range(lm_layers)
        ]

    return batch_states