def get_all_logits(predictions, features):
    """
    Check assertions for predictions ans features length.
    If passes, return all start/end logits.

    :param predictions: predictions
    :param features: tokenized and divided contexts(The processed dataset, cut by max_sequence_length)
    :return: all start/end logits if pass assertions.
    """
    assert len(predictions) == 2, "`predictions` should be a tuple with two elements (start_logits, end_logits)."
    all_start_logits, all_end_logits = predictions

    assert len(predictions[0]) == len(features), f"Got {len(predictions[0])} predictions and {len(features)} features."

    return all_start_logits, all_end_logits