"tf_device.launch"() ( {
        "tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
        tf_device.return
      }) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()
      %1 = "tf_device.launch"() ( {
        %2 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<i8>, tensor<3x!tf_type.string>) -> tensor<i8>

   // CHECK: _TPUCompileMlir
   // CHECK-SAME: Z
   // CHECK: tf_device.launch
   // CHECK-NEXT: _TPUCompileMlir
   // CHECK-SAME: X
   // CHECK: tf_device.launch
   // CHECK-NEXT: _TPUCompileMlir
   // CHECK-SAME: Y
   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()
   "tf_device.launch"() ({

          %b1 = "tf.AddV2"(%barg0, %0) {T = i32, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
          // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
          // CHECK-NEXT: "tf._TPUCompileMlir"()
          %compile:2 = "tf_device.launch"() ({
            %b2:2 = "tf._TPUCompileMlir"() {
              NumDynamicShapes = 0 : i64,
              // The metadata encodes 2 parameter and 2 return values.

     // CHECK:      %[[RESOURCE_OUT:.*]] = "tf_device.launch"()
     // CHECK-SAME: TPU_REPLICATED_CORE_0
     // CHECK-NEXT:   %[[READ_OUT:.*]] = "tf.ReadVariableOp"(%[[RI_0]])
     // CHECK-NEXT:   tf_device.return %[[READ_OUT]]
     %0 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource<tensor<4xf32>>>) -> tensor<4xf32>
     %1 = "tf.A"() : () -> (tensor<2x!tf_type.string>)
     "tf_device.launch"() ({

      return %cluster : tensor<i32>
    ```

    becomes:

    ```mlir
    %0 = "tf_device.launch"() ( {
      %3 = "tf.A"(%arg0) : (tensor<i32>) -> tensor<i32>
      tf_device.return %3 : tensor<i32>
    }) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
    %1 = "tf_device.cluster"() ( {

        %2 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

        // Note that tf.C is moved before tf_device.launch.
        // CHECK: %[[C_OUTPUT:[0-9]*]] = "tf.C"(%[[ARG_0]])

        // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
        // CHECK: <{device = "tpu0"}>
        // CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[A_OUTPUT]]) : (tensor<?xi32>) -> tensor<?xi32>

    after a _TPUCompileMlir op if there is no predecessor _TPUCompileMlir op.
    Both the embedding op and compile op are assumed to be wrapped in separate
    tf_device.launch() ops. This is because the embedding op is head outside
    compiled and the compile op is wrapped in launch to execute on host
    during TPURewritePass.

    For example, the tf.OpA with the `mini_batch_splits` attribute will be

     * Returns the extra arguments to use to launch the JVM for the process. Does not include system properties and the
     * minimum/maximum heap size.
     *
     * @return The immutable list of arguments. Returns an empty list if there are no arguments.
     */
    @Nullable @Optional @Input
    List<String> getJvmArgs();

    /**
     * Sets the extra arguments to use to launch the JVM for the process. System properties

    // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:3 = "tf_device.launch"
    // CHECK-NEXT: "tf._TPUCompileMlir"()
    // CHECK: "tf_device.launch"
    // CHECK: "tf.TPUCompileSucceededAssert"(%[[COMPILE_OUTPUT]]#0)
    // CHECK: [[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:2 = "tf_device.parallel_execute"
    // CHECK: "tf_device.launch"() <{device = "/job:worker/replica:0/task:0/device:TPU:0"}>

import org.gradle.launcher.daemon.server.api.DaemonCommandAction;
import org.gradle.launcher.daemon.server.api.HandleInvalidateVirtualFileSystem;
import org.gradle.launcher.daemon.server.api.HandleReportStatus;
import org.gradle.launcher.daemon.server.api.HandleStop;
import org.gradle.launcher.daemon.server.exec.CleanUpVirtualFileSystemAfterBuild;
import org.gradle.launcher.daemon.server.exec.DaemonCommandExecuter;