}
        ops_to_erase.insert(input_op);
      }
    }
    for (auto op : ops_to_erase) op->erase();
    return LogicalResult::success();
  }

  LogicalResult PatchCollectiveGatherInstanceKey(func::FuncOp func) {
    // We're expecting the original model to have a single CollectiveGatherV2Op
    // that gets split into 3 copies in the start_step_0, start_step_1 and