func.return %arg0 : tensor<i32>
}

func.func @testClusterFuncOpFails(%arg0: tensor<i32>) -> tensor<i32> {
   // expected-error@below {{failed TF functional to executor validation, op tf_device.cluster_func is not allowed}}
  %cluster = "tf_device.cluster_func"(%arg0) {func = @_func} : (tensor<i32>) -> tensor<i32>
 func.return %cluster : tensor<i32>
}

// -----

// CHECK-LABEL: func @testTFDialect

//   RemoveSingletonParallelExecuteOp can be removed.

// Wrap `cluster_func` in a `ParallelExecute` with only one child. This
// can be used to canonicalize IR, so there is always one `ParallelExecute`.
tf_device::ParallelExecuteOp BuildParallelExecuteOp(
    tf_device::ClusterFuncOp cluster_func, OpBuilder* builder);

// Unwrap `parallel_execute`'s contents if it only has one child.

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({
    // CHECK: %compilation_status, %program = "tf._TPUCompileMlir"()

    %5:2 = tf_device.replicate([%0, %arg0] as %ri_0: tensor<?xi32>) {n = 2 : i32} {
      %2 = "tf._F"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
      %3 = "tf.Identity"(%1) : (tensor<?xi32>) -> tensor<?xi32>
      %4 = "tf_device.cluster_func"(%ri_0, %3, %2) {func = @_func, device = ""} : (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
      tf_device.return %4 : tensor<?xi32>
    }
    %6 = "tf._C"(%5#1) : (tensor<?xi32>) -> tensor<?xi32>

          VerifyInputDialectToExecutorPass> {
 public:
  void runOnOperation() override;
};

bool IsTfDeviceClusterFuncOp(Operation* op) {
  std::string kClusterFuncOpName = "tf_device.cluster_func";
  return op->getName().getStringRef().str() == kClusterFuncOpName;
}

void VerifyInputDialectToExecutorPass::runOnOperation() {
  Operation* func_op = getOperation();

  func.func @missing_num_cores_per_replica() {
    // expected-error@+1 {{requires attribute 'num_cores_per_replica'}}

    Bubbles up sharding configuration from `cluster_func` regions into
    the attributes of `cluster_func`. This is done by parsing the
    `XlaSharding` / `TPUPartitionedOutput` / `TPUPartitionedInput` ops inside
    `cluster_func`.

    For example, given the following `cluster_func` wrapping `func`:

    ```mlir
      func @test(%arg0: tensor<*xi32>) {
        "tf_device.cluster_func"(%arg0) {
            func = @func,

}

int MovePreservedParallelExecuteChildren(
    int num_cores_per_replica,
    llvm::SmallVector<mlir::Type, 8>& concatenated_output_types,
    mlir::OpBuilder* builder, mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::tf_device::ParallelExecuteOp old_parallel_execute,
    mlir::tf_device::ParallelExecuteOp* new_parallel_execute) {
  // `num_moved_children` is the number of children that will be preserved.

  func.func @main() {

  func.func @main(%arg0: tensor<*xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) {