brevity)
```mlir
  func.func @remove_first_result(%arg0, %arg1) {
    %0:2 = "tf.WhileRegion"(%arg0, %arg1) ({
    ^bb0(%arg2, %arg3):
      %1 = "tf.OpA"() {is_stateless = true}
      "tf.Yield"(%1)
    }, {
    ^bb0(%arg2, %arg3):
      %1 = "tf.OpB"(%arg2) {is_stateless = true}
      %2 = "tf.OpC"(%arg3) {is_stateless = true}
      "tf.Yield"(%1, %2)
    }) {is_stateless = true}

      %w1, %w2 = "tf.WhileRegion"(%a, %b) ({
      ^bb0(%arg1: tensor<i32>, %arg2: tensor<f32>):
        %7 = "tf.H"(%arg1) :  (tensor<i32>) -> tensor<i1>
        "tf.Yield"(%7) : (tensor<i1>) -> ()
      }, {
      ^bb0(%arg1: tensor<i32>, %arg2: tensor<f32>):
        %8 = "tf.C"(%arg1) : (tensor<i32>) -> tensor<i32>
        %9 = "tf.D"(%arg1, %arg2) : (tensor<i32>, tensor<f32>) -> tensor<f32>

  func.func private @testIfElse(%arg0:tensor<*xf32>, %arg1:tensor<*x!tf_type.resource>) -> tensor<*xf32> {
    %val = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource>) -> (tensor<0xf32>)
    %0 = "tf.Mul"(%val, %arg0) : (tensor<0xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  func.func @f(%arg0: tensor<i1>, %arg1: tensor<2xf32>) -> tensor<2xf32> {

}

func.func @while_body(%arg0: tensor<*xi32>) -> (tensor<*xi32>) {
  %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>

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>

    func.return %1: tensor<*xf32>
  }
  func.func private @composite_matmul_fn(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x1024xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

    %1:2 = "tfl.while"(%cst, %0) ({
    ^bb0(%arg1: tensor<*xi32>, %arg2: tensor<*x!tf_type.resource<tensor<*xi32>>>):
      %2 = "tf.ReadVariableOp"(%arg2) {dtype = i32} : (tensor<*x!tf_type.resource<tensor<*xi32>>>) -> tensor<*xi32>
      %3 = "tfl.greater"(%arg1, %2) : (tensor<*xi32>, tensor<*xi32>) -> tensor<i1>
      "tfl.yield"(%3) : (tensor<i1>) -> ()
  },  {
    ^bb0(%arg3: tensor<*xi32>, %arg4: tensor<i32>):

// CHECK: %[[PARTITIONEDCALL:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST]])
// CHECK-SAME: f = @composite_matmul_fn_1}>
// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[UNQUANTIZED_OUTPUT:.*]] = "tf.MatMul"(%arg0, %arg0)
// CHECK: }

// CHECK-LABEL: private @composite_matmul_fn_1
// CHECK-NEXT: %[[OUT:.*]] = "tf.MatMul"(%arg0, %arg1)
// CHECK-NEXT: return %[[OUT]]
}

// -----

    // non_tpu_ops
    %res_n = "tf.Add"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>

    return %res_n : tensor<i32>
  }
  func.func private @while_cond(%arg0: tensor<i32>) -> tensor<i1> {
    %0 = "tf.Less"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    return %0 : tensor<i1>
  }

func.func @main(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
  %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
  func.return %0:  tensor<4x10x17xf32>

// CHECK-LABEL: main