// Start step 1
C_1 = cond(args_1)
N_1 = non_tpu(args_1)
if (C_1) {
   F_1 = forward(args_1, N_1)
   // T_1 = core_tpu() is not evaluated here.
   // B_1 = backward() is not evaluated here.
}

// Partial update of args. We expect this to be sufficient
// for evaluating cond().
args_2a = update_args(args_1, N_1)  // NO T_1 here

// Conditional for step 2
C_2 = cond(args_2)

new_while_body (new_args) {  // starts at i==2