%3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]])
// CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[matmul]])

// RUN: tf-opt %s -tf-tpu-identity-pruning | FileCheck %s --dump-input=always

// Tests Identity op in cluster is pruned away.

// CHECK-LABEL: func @testIdentity
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @testIdentity(%arg0: tensor<i32>) {
  // CHECK-NOT:  "tf.Identity"
  // CHECK:      "tf_device.cluster"
  // CHECK-NEXT: tf_device.return [[ARG0]]
  %0 = "tf_device.cluster"() ({
    %1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>

  // CHECK-NEXT: bb0(%[[ARG_3:.*]]: tensor<i32>, %[[ARG_4:.*]]: tensor<*xf32>)
  ^bb0(%arg3: tensor<i32>, %arg4: tensor<*xf32>):
    // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.Identity"(%[[ARG_3]])
    %1 = "tf.Identity"(%arg3) : (tensor<i32>) -> tensor<i32>
    %2 = "tf.Identity"(%arg1) : (tensor<i32>) -> tensor<i32>
    // CHECK-NEXT: %[[RESULT_2:.*]] = "tf.NotEqual"(%[[RESULT_1]], %[[ARG_1]])

    // CHECK:      tf.Identity
    // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:TPU:0"
    %1:2 = tf_executor.island wraps "tf.Identity"(%0#0) : (tensor<i64>) -> tensor<i64>
    // CHECK:      tf.IdentityN
    // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:TPU:0"

// CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[ARG0]], %[[cast]])
// CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[conv]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
// CHECK: return %[[identity]] : tensor<1x3x2x2xf32>

// -----

// Tests that `tf.BatchFunction` is inlined.

Harshavardhana <******@****.***> 1627436271 -0700

jiuker <******@****.***> 1684487613 +0800

  func.return
}
// CHECK-LABEL: empty_island_multiple_data_results
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<*xf32>, %[[ARG_1:.*]]: tensor<*xi32>)
// CHECK: %[[IDENTITY_N:.*]]:2 = "tf.IdentityN"
// CHECK-SAME: (%[[ARG_0]], %[[ARG_1]])
// CHECK: tf_executor.yield %[[IDENTITY_N]]#0, %[[IDENTITY_N]]#1

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {

func.func @while_main(%arg0: tensor<?x256x256xf32>) -> (tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "Identity,Identity_1,Identity_2"}} {
  %cst = arith.constant dense<1.000000e+00> : tensor<256x256xf32>
  %cst_0 = arith.constant dense<0> : tensor<i32>
  %cst_1 = arith.constant dense<-1> : tensor<i32>

    return absl::OkStatus();
  }

  // Create a single IdentityN that is dead if and only if
  // `predicate_as_control` is dead.
  //
  // IdentityN is also special in that, unlike `Identity`, it does not place
  // int32 inputs in host memory.  Placing int32 inputs in host memory would
  // defeat the purpose of adding this indirection.
  ops::IdentityN identity_n(root.WithOpName("int32_id_n"), int32_inputs);