def METHOD_NAME(before, after, length):
    # Make sure all the batch tensors have the same shape
    input_ids_after_shape = after['input_ids'].shape
    # Just make sure the lengths are correct
    for k in before.keys():
        assert k in after, 'No keys should be removed during sequence reshaping.'
        assert after[
            k].shape == input_ids_after_shape, 'All tensors should have the same size after sequence reshaping.'
        b_numel = before[k].shape[0] * before[k].shape[1]
        a_numel = after[k].shape[0] * after[k].shape[1]
        assert a_numel >= b_numel - length, 'Sequence reshaping should throw away at most curr_sequence_length tokens.'
        import torch
        assert torch.all(after[k][0] == before[k][
            0, :input_ids_after_shape[1]]), 'Sequence reshaping should not change the token order.'
    for k in after.keys():
        assert k in before, 'No keys should be added during sequence reshaping.'