def trace_overall_index(batch_idx, test_time_indices):
    """
    Given
        batch_idx: the particular program this all corresponds to
        test_time_indices: list of which (tests, timestep) to select

    Returns (trace_idxs, time_idxs)
        trace_idxs: the indices int the traces to be returned
        time_idxs: the indices into which timesteps should be returned
    """
    # this assumes 5 tests exactly
    assert all(test < 5 for test, _ in test_time_indices)
    return [batch_idx * 5 + test for test, _ in test_time_indices], [time for _, time in test_time_indices]