def METHOD_NAME(self, state: Any, i: int, new_id: int = None) -> Any:
    """Select state with relative ids in the main beam search.
    Args:
        state: Decoder state for prefix tokens
        i (int): Index to select a state in the main beam search
        new_id (int): New label index to select a state if necessary
    Returns:
        state: pruned state
    """
    return None if state is None else state[i]