%add = "tf.AddV2"(%call, %call) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
    "tf.Yield"(%call, %add) : (tensor<4xf32>, tensor<4xf32>) -> ()
  },  {
    %call_0 = func.call @test_if_else1(%arg1) : (tensor<4xf32>) -> tensor<4xf32>
    "tf.Yield"(%call_0, %call_0) : (tensor<4xf32>, tensor<4xf32>) -> ()
  }) {is_stateless = false} : (tensor<i1>) -> (tensor<*xf32>, tensor<4xf32>)

module {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %1 = "tfl.mul"(%0, %arg2) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

      %outputs_9, %control_10 = tf_executor.island wraps "tf.Mul"(%outputs_7, %outputs) {T = "tfdtype$DT_FLOAT", _tpu_replicate = "cluster", device = "", name = "mul"} : (tensor<0xf32>, tensor<f32>) -> tensor<0xf32>
      %outputs_11, %control_12 = tf_executor.island wraps "tf.Placeholder"() {device = "", dtype = "tfdtype$DT_FLOAT", name = "y", shape = "tfshape$dim { }"} : () -> tensor<0xf32>

  func.return %1 : tensor<1xf32>
}

func.func @cond_true(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

func.func @cond_false(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {

func.func @main(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
  %1 = "tfl.squared_difference"(%arg0, %0) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("squared_difference")
  %2 = "tfl.mul"(%0, %1) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("mul")

module {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
  %0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.rsqrt", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xf32>) -> tensor<2xf32>
  func.return %0 : tensor<2xf32>
}
}

// CHECK:       module
// CHECK-NEXT:  func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NEXT:  %0 = stablehlo.rsqrt %arg0 : tensor<2xf32>
// CHECK-NEXT:  return %0 : tensor<2xf32>

    //body
    ^bb0(%arg_body: tensor<1xf32>):
    %result_body = "tfl.add"(%arg_body, %arg_body) { fused_activation_function = "NONE" } : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
    "tfl.yield"(%result_body) : (tensor<1xf32>) -> ()
  }) : (tensor<1xf32>) -> (tensor<1xf32>)
  %tmp5 = "tfl.add"(%tmp4, %tmp2) { fused_activation_function = "NONE" } : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

// CHECK-NEXT:   "tfl.yield"(%[[MY_CUSTOM]]) : (tensor<4xf32>) -> ()
// CHECK-NEXT: }) {fused_activation_function = "RELU", int_attr = 2 : i32} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK-NEXT: %[[CUSTOM_2:.*]] = "tfl.custom_tf"(%[[MUL]], %[[CST]]) ({
// CHECK-NEXT: ^bb0(%arg1: tensor<4xf32>, %arg2: tensor<4xf32>):

  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) {
    device = "", narrow_range = false, num_bits = 8 : i64
  } : (tensor<3xf32>, tensor<f32>, tensor<f32>) -> tensor<3xf32>
  %1 = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg1, %cst_1, %cst_2) {
    device = "", narrow_range = true, num_bits = 8 : i64
  } : (tensor<4x3xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<4x3xf32>
  func.return %0, %1 : tensor<3xf32>, tensor<4x3xf32>
}

// CHECK-SAME: (%[[ARG0:.*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)
func.func @testReadVariableOpColocated(%arg0: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
  // CHECK:      tf_device.replicate
  // CHECK-SAME: (%[[ARG0]] as %[[RI_0:[a-z0-9]*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)
  tf_device.replicate(%arg0 as %arg1: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
    _mirrored_variable_indices = [0],