def extract_iteration()

in visualize/generate.py [0:0]


def extract_iteration(epoch_data):
    data_by_iter = []
    current_iter = []
    broadcast_count = 0

    for item in epoch_data:
        if item.comm_type == CommType.epoch_end:
            if current_iter:
                data_by_iter.append(current_iter)
            break

        if broadcast_count < 2:
            if item.comm_type == CommType.broadcast:
                broadcast_count += 1
            current_iter.append(item)
        else:
            if item.comm_type == CommType.broadcast:
                data_by_iter.append(current_iter)
                current_iter = [item]
                broadcast_count = 1
            else:
                current_iter.append(item)



    return data_by_iter