%1 = "tf.Add"(%val_var1, %cst_0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %2 = "tf.Mul"(%1, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %3 = "tf.AddV2"(%arg2, %1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %4 = "tf.Div"(%arg3, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "tf.Yield"(%3, %4) : (tensor<i32>, tensor<i32>) -> ()

    tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
  }
  return %0#0, %0#1 : tensor<i32>, tensor<i32>
}

// -----

func.func @validate_tpu_replicate_wrong_attr(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  %0:2 = tf_executor.graph {

  // CHECK: [[z:%.*]] = tf_mlrt.executeop
  %z = "tf.Identity"(%x) {__op_key = 1: i32}: (tensor<i32>) -> tensor<i32>
  // CHECK: [[y:%.*]] = tf_mlrt.await [[y_future]]
  // CHECK-NEXT: tf_mlrt.executeop([[y]]
  %w = "tf.AddV2"(%y, %z) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK-NEXT: tf_mlrt.executeop([[y]]

  %a6 = "tf.AddV2"(%a5, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %a7 = "tf.AddV2"(%a6, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>

  %b0 = "tf.Sub"(%a3, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>

  %1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
  func.return %1 : tensor<i1>
}

func.func @side_effect_body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
  %0 = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %1 = "tf.AddV2"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>

}

func.func @consecutive_add_sub(%arg0: tensor<i32>) -> (tensor<i32>) {
  %cst = "tf.Const"() {value = dense<-18> : tensor<i32>} : () -> tensor<i32>
  %cst_1 = "tf.Const"() {value = dense<-12> : tensor<i32>} : () -> tensor<i32>
  %0 = "tf.AddV2"(%arg0, %cst) {T = i32, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %1 = "tf.Sub"(%0, %cst_1) {T = i32, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  return %1 : tensor<i32>

  %1 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee1} : (tensor<i32>) -> tensor<i32>
  func.return %1 : tensor<i32>
}

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

    %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %cst_2 = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
    %0:2 = "tf.While"(%cst_1, %cst_2) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {

      %a = "tf.A"() {_xla_outside_compilation = "cluster1"} : () -> tensor<i32>
      %b = "tf.B"(%a) : (tensor<i32>) -> tensor<i32>
      %c = "tf.C"(%b) : (tensor<i32>) -> tensor<i32>
      tf_device.return %a, %c, %b : tensor<i32>, tensor<i32>, tensor<i32>
    }) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} : () -> (tensor<i32>, tensor<i32>, tensor<i32>)

// CHECK: "tf.opA"
// CHECK-NOT: island
// CHECK: name = "A"
// CHECK-NOT: island
// CHECK: name = "C"
// CHECK-NOT: island
// CHECK: name = "E"
// CHECK: island {{.*}}name = "B"
// CHECK: island {{.*}}name = "D"

    tf_executor.fetch %1#0, %2#0, %3#0, %4#0, %5#0, %6#0 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
  }
  func.return %0#0 : tensor<i32>
}