mlir::TF::SplitIsland(curr_island, control_type);
        });
  }
}

}  // anonymous namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateLaunchToDeviceAttributePass(
    bool legacy_graph_export) {
  return std::make_unique<LaunchToDeviceAttributePass>(legacy_graph_export);
}

}  // namespace TFDevice

  pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateParallelExecuteToIslandsPass(
      /*legacy_graph_export=*/false));
  pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass(
      /*legacy_graph_export=*/false));

  // Do a single pass to encode necessary control deps in the IR according to
  // the results of side effect analysis.
  pm.addPass(

  add_pass(mlir::TFDevice::CreateReplicaIDToDeviceOrdinalPass());
  add_pass(mlir::TFDevice::CreateParallelExecuteToIslandsPass(
      /*legacy_graph_export=*/true));
  add_pass(mlir::TFDevice::CreateLaunchToDeviceAttributePass(
      /*legacy_graph_export=*/true));
  pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUDevicePropagationPass());
  pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUColocateSplitsPass());