def rnn_state_trim(rnn_type, states):
    """
    Trims a given state depending on the type of RNN.

    It mostly affects LSTMs, which return a tuple as a state, by only returning
    the hidden vector.

    Args:
        rnn_type (str): type of RNN (e.g. "lstm").
        states: the states returned by the RNN.
          [2, batch_size, size] for LSTM
          [batch_size, size] for others

    Returns:
        states: the trimmed states.
          [batch_size, size]
    """
    if rnn_type == "lstm":
        return states[0]

    return states