import torch


def unpack_bidirectional_lstm_state(state, num_directions=2):
    """
    Unpack the packed hidden state of a BiLSTM s.t. the first dimension equals to the number of layers multiplied by
    the number of directions.
    """
    batch_size = state.size(1)
    new_hidden_dim = int(state.size(2) / num_directions)
    return torch.stack(torch.split(state, new_hidden_dim, dim=2), dim=1).view(-1, batch_size, new_hidden_dim)