}

  // Replace original terminator with new terminator for returning merged
  // execute results.
  Operation* old_terminator = execute_region->front().getTerminator();
  builder->setInsertionPointToEnd(&execute_region->front());
  builder->create<tf_device::ReturnOp>(old_terminator->getLoc(),
                                       merged_execute_launch.getResults());
  old_terminator->erase();

        island_op.getLoc(), llvm::ArrayRef<Type>{},
        island_op.getControl().getType(), island_operands);
    island_sink.getBody().push_back(new Block);
    builder.setInsertionPointToEnd(&island_sink.GetBody());
    builder.create<tf_executor::YieldOp>(island_op.getLoc(),
                                         llvm::ArrayRef<Value>{});

    for (auto i : return_indices[index_and_value.index()]) {
      merged_then_yield_values.push_back(
          if_op.getThenBranch().front().getTerminator()->getOperand(i));
    }
  }
  builder.setInsertionPointToEnd(&new_if_op.getThenBranch().front());
  builder.create<TF::YieldOp>(
      first_if.getThenBranch().front().getTerminator()->getLoc(),
      /*operands=*/merged_then_yield_values);

        island_op.getLoc(), llvm::ArrayRef<Type>{},
        tf_executor::ControlType::get(island_op.getContext()), island_operands);
    island_sink.getBody().push_back(new Block);
    builder.setInsertionPointToEnd(&island_sink.GetBody());
    builder.create<tf_executor::YieldOp>(island_op.getLoc(),
                                         llvm::ArrayRef<Value>{});

      loc, result_types, control_type, control_inputs);
  island.getBody().push_back(new Block);
  Block* block = &island.getBody().back();
  OpBuilder island_builder(original_island);
  island_builder.setInsertionPointToEnd(block);
  if (sub_op) {
    sub_op->replaceAllUsesWith(island.getOutputs());
    sub_op->moveBefore(block, block->begin());
    island_builder.create<tf_executor::YieldOp>(loc, sub_op->getResults());
  } else {

  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;
}

// Creates a new_launch after `before_op` with `ops_to_move` from

        if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
          inner_op->removeAttr(kDeviceAttr);
        }
      }
    });
  }

  // Add terminator.
  builder.setInsertionPointToEnd(body);
  builder.create<mlir::tf_device::ReturnOp>(last_cluster_op->getLoc(), results);

  // Replaces uses of cluster ops results outside of cluster with the associated
  // `tf_device.cluster` results.

  // (operands) and 'outputs' (results).
  OpBuilder builder(module);

  std::vector<Type> input_types = GetValueTypes(inputs);
  std::vector<Type> output_types = GetValueTypes(outputs);
  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

    return op->getResults();
  }
  mlir::Region region;
  region.push_back(new mlir::Block);
  auto saved_pos = op_builder.saveInsertionPoint();
  op_builder.setInsertionPointToEnd(&region.front());
  mlir::Operation* cloned_op = op_builder.clone(*op);
  // Add the yield operation.
  op_builder.create<mlir::TFL::YieldOp>(op_loc, cloned_op->getResults());

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

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

  return launch;
}

// Checks if an operation is a supported TPU embedding op.