def _get_aiob_compute_time()

in workload_generator/AIOB_simAI_workload_generator.py [0:0]

• 16 lines of code
• 6 McCabe index (conditional complexity)

def _get_aiob_compute_time(compute_cache, forward_or_backward, stage):
    compute_time_map = compute_cache
    if stage == "grad":
        prefix = stage + "_" + forward_or_backward
    elif stage == "embedding":
        prefix = "Emb"
    elif stage == "final":
        prefix = "attention" + "_" + forward_or_backward
    else:
        prefix = stage + "_" + forward_or_backward

    for key, value in compute_time_map.items():
        if prefix == key:

            compute_time = compute_time_map.get(key)
            return compute_time

    print("[warn] can't match any stage", stage)
    return 1