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());

  // come before TPUClusterCleanupAttributes pass or else the device attribute
  // will be removed from launch causing an error.
  pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass());

  // Note that the region-based control-flow produced here still contains
  // function call ops which get inlined by the subsequent inliner pass.

           "islands created by this pass should be split after the replicated "
           "ops have been lifted.">
  ];

  let constructor = "TFDevice::CreateLaunchToDeviceAttributePass()";
}

def XlaInlineDeviceOpsPass : Pass<"tf-xla-inline-device-ops", "ModuleOp"> {
  let summary = "Inline all Cluster op based in the parent region";

// attribute to each TensorFlow dialect op in the body based on the `device`
// attribute on the `tf_device.launch`.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLaunchToDeviceAttributePass(
    bool legacy_graph_export = true);

// Creates a pass to ensure that the `_xla_outside_compilation` and
// tf_device.launch op no longer exist after Outside Compilation is complete.