%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]]

}

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>

    %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>)

    %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %1 = "tf.AddV2"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %2 = "tf.Div"(%arg1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    func.return %1, %2 : tensor<i32>, tensor<i32>
  }

      %outputs_6, %control_7 = tf_executor.island wraps "tf.Sub"(%outputs_0, %arg0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
      %outputs_8, %control_9 = tf_executor.island wraps "tf.FloorDiv"(%outputs_6, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
      tf_executor.fetch %outputs_8 : tensor<i32>
    }
    func.return %0 : tensor<i32>
  }
}

  func.func @convert_cluster_func(%arg0: tensor<i32>) -> tensor<i32> {
    // CHECK: "tf.XlaLaunch"(%arg0) <{function = @func, operandSegmentSizes = array<i32: 0, 1, 0>}> : (tensor<i32>) -> tensor<i32>
    %0 = "tf_device.cluster_func"(%arg0) {func = @func} : (tensor<i32>) -> tensor<i32>
    func.return %0 : tensor<i32>
  }

  func.func @func(%arg0: tensor<i32>) -> tensor<i32> {
    func.return %arg0 : tensor<i32>
  }
}

// -----

  %1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %2 = "tf.Add"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  "tf.AssignVariableOp"(%arg, %2) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  func.return %arg : tensor<!tf_type.resource<tensor<i32>>>
}