scope.graph()->AddEdge(u.node(), 0, launch, 4);
  scope.graph()->AddEdge(v.node(), 0, launch, 5);
  scope.graph()->AddEdge(w.node(), 0, launch, 6);

  auto out0 =
      ops::XlaClusterOutput(scope.WithOpName("Out0"), Output(launch, 0));
  auto out1 =
      ops::XlaClusterOutput(scope.WithOpName("Out1"), Output(launch, 1));
  auto out2 =
      ops::XlaClusterOutput(scope.WithOpName("Out2"), Output(launch, 2));

// CHECK-NEXT:   tf_device.return %[[OP_A]]
// CHECK:      %[[LAUNCH_B:[0-9]*]] = "tf_device.launch"
// CHECK-SAME: device = "b"
// CHECK-NEXT:   %[[OP_B:[0-9]*]] = "tf.opB"(%[[SHAPE]], %[[LAUNCH_A]])
// CHECK-NEXT:   tf_device.return %[[OP_B]]
// CHECK: tf_device.replicate([{{.*}}] as %[[RI:[a-z0-9]+]]: tensor<*xf32>)
// CHECK-NEXT:   %[[LAUNCH_C:[0-9]*]] = "tf_device.launch"
// CHECK-SAME:   device = "c"

  builder->setInsertionPointAfter(before_op);
  auto launch = builder->create<tf_device::LaunchOp>(
      before_op->getLoc(), builder->getStringAttr(host_device),
      launch_result_types);
  launch.getBody().push_back(launch_block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(before_op->getLoc(), launch_results);

  return launch;
}

    } else {
      break;
    }
  }
  return false;
}

// Wraps block in a Launch. External uses of ops in the block will be return
// values of the Launch and remapped to the Launch results. If `before` is set
// to true, the Launch is created before `op`. Otherwise the Launch is created
// after `op`.
mlir::tf_device::LaunchOp CreateLaunchForBlock(OpBuilder* builder,

replica, the execution semantics follow standard sequential behavior. Ops in the
tf_device.replicate wrapped with a tf_device.launch will have its device set to
the associated replicated device from `devices` if the tf_device.launch refers
to an aliased device name. Otherwise the device already set in tf_device.launch
is used instead.

Operands are replicated inputs and packed inputs.

}

// Wraps single op in `tf_device.launch` for explicit device assignment.
void WrapOpInLaunch(OpBuilder* builder, Location loc, Operation* op,
                    llvm::StringRef device) {
  OpBuilder::InsertPoint insert_point = builder->saveInsertionPoint();

  auto launch = builder->create<tf_device::LaunchOp>(
      loc, builder->getStringAttr(device), op->getResultTypes());
  launch.getBody().push_back(new Block);

          : execute_launch.getResults();
  for (auto execute_output : llvm::enumerate(execute_outputs)) {
    // TODO(lyandy): Handle updates to resource writes by remapping to parent
    // launch result and checking if launch result is an AssignVariableOp.
    auto result = execute_output.value();
    if (!result.hasOneUse()) {
      if (VLOG_IS_ON(2)) {
        bool any_user_is_assign = false;

def LaunchToDeviceAttributePass : Pass<"tf-launch-to-device-attribute", "mlir::func::FuncOp"> {
  let summary = "Hoists and annotates device launch inner ops with associated device attribute.";

  let description = [{
    This pass hoists a `tf_device.launch` body and assigns a `device` attribute
    to each TensorFlow dialect op in the body based on the `device` attribute on
    the `tf_device.launch`. If a TensorFlow dialect op already has a device

        when:
        withConnection { connection ->
            TestLauncher launcher = connection.newTestLauncher().withTests(descriptor)
            collectOutputs(launcher)
            launcher.run()
        }

        then:
        outputContains("BUILD SUCCESSFUL")
    }

    def "can launch test with test launcher via test filter targeting a specific task"() {
        setup:

}

//===----------------------------------------------------------------------===//
// tf_device.launch
//===----------------------------------------------------------------------===//

// Checks if a tf_device.launch wraps a single operation and the single
// operation results are perfectly forwarded to the launch return.
bool LaunchOp::WrapsSingleOp() { return BlockWrapsSingleOp(&GetBody()); }