func.func @stateful_composite_op_control_else(%arg0: tensor<*x!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
  %0 = tf_executor.graph {
    %outputs, %control = tf_executor.island {
      %1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
      "tf.AssignVariableOp"(%arg0, %1) : (tensor<*x!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
      tf_executor.yield %1 : tensor<i32>
    }
    tf_executor.fetch %outputs : tensor<i32>
  }

    %4:2 = tf_executor.island wraps "tf.Mul"(%2#1, %2#1) {T = "tfdtype$DT_INT32"} : (tensor<i32>, tensor<i32>) -> tensor<i32> loc("Multiplication")
    %5:3 = tf_executor.Merge %3#0, %4#0 : tensor<i32> {device = "", N = 2, T = "tfdtype$DT_INT32"} loc("Merge")
    tf_executor.fetch
  }
  func.return
}

  %0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
  func.return %0 : tensor<256x30x30x16xf32>
}

  func.return %arg0 : tensor<i32>
}

// -----

// CHECK-LABEL: func @stateful_partitioned_call_multi_in_out
// CHECK-SAME: ([[ARG0:%.+]]: tensor<i32>, [[ARG1:%.+]]: tensor<i32>)
func.func @stateful_partitioned_call_multi_in_out(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {

// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[DOT_GENERAL_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>)

      func.func @main() -> tensor<!tf_type.variant<tensor<64x1xbf16>>> {
      // unknown rank
      %elem_shape = "tf.Const"() <{value = dense<-1> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
      // zero reserved elements

  }) {is_stateless = false} : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
  func.return %0#1 : tensor<1xf32>
}

func.func @cond(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>) -> tensor<i1> {
  %cst = arith.constant dense<0> : tensor<i32> loc("Const")
  %0 = "tfl.greater"(%arg0, %cst) : (tensor<*xi32>, tensor<i32>) -> tensor<i1>
  func.return %0 : tensor<i1>
}

// CKECK-SAME: <{calibration_method = 5 : i32, id = "0", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"

  }

  func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
    %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
    %1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
    func.return %1 : tensor<i1>
  }

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

  %0 = "tf.Rank"(%arg0) : (tensor<*xi32>) -> tensor<i32>
  %1 = "tf.Add"(%0, %arg0): (tensor<i32>, tensor<*xi32>) -> tensor<*xi32>
  func.return %1: tensor<*xi32>
}

func.func @while_cond(%arg0: tensor<*xi32>) -> tensor<*xi1> {
  %cst = arith.constant dense<10> : tensor<i32>
  %0 = "tf.Less"(%arg0, %cst) {T = i32, device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
  func.return %0 : tensor<*xi1>
}