%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
  %bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
  %output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

  %0 = "tfl.pseudo_const"() {value = dense<[2, 2, 81]> : tensor<3xi32>} : () -> tensor<3xi32>
  %1 = "tfl.pseudo_const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
  %2:3 = "tfl.split_v"(%arg0, %0, %1) {num_splits = 3 : i32} : (tensor<?x13x13x85xf32>, tensor<3xi32>, tensor<i32>) -> (tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>)

  %stride = "tf.Add"(%arg2, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %cst1 = arith.constant dense<44> : tensor<i32>
  %limit = "tf.Add"(%arg2, %cst1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: [[Result:%.*]]:4 = "tf.While"(%arg0, %arg1, %{{.+}}, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) (
    {
      ^bb0(%carg0: tensor<*xf32>, %carg1: tensor<i32>):

  %conv = "tfl.conv_2d"(%arg0, %w, %b) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32} : (tensor<1x224x224x3xf32>, tensor<64x3x3x3xf32>, tensor<64xf32>) -> tensor<1x112x112x64xf32>

    },
    {
      // loop body
      ^bb0(%barg0: tensor<i32>, %barg1: tensor<i32>):
      %one = arith.constant dense<1> : tensor<i32>
      %sub0 = "tf.Sub"(%barg0, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      %sub1 = "tf.Sub"(%barg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      "tf.Yield"(%sub0, %sub1) : (tensor<i32>, tensor<i32>) -> ()
    }

    %cst = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
    %11:4 = "tf.Split"(%cst, %0) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
    %cst_0 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

         %cond = "tf.NotEqual"(%carg0, %limit) : (tensor<i32>, tensor<i32>) -> tensor<i1>
         "tf.Yield"(%cond) : (tensor<i1>) -> ()
      }, {
       // CHECK: {{.*}}({{.+}}: tensor<i32>, {{.+}}: tensor<1x2x3xf32>):
       ^bb0(%barg0: tensor<i32>, %barg1: tensor<*xf32>):
        %one = arith.constant dense<1> : tensor<i32>
        %sub = "tf.Sub"(%barg0, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>

  // CHECK: func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<3x11xf32>
  // CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<3x11xf32>
  // CHECK: return %[[CONSTANT_1:.*]]
  // CHECK: }

  // CHECK: func private @_stablehlo_main_0
  // CHECK-SAME: (%arg0: tensor<i32>, %[[INPUT_1:.*]]: tensor<3x11xf32>)

  %0 = "tfl.while"(%arg0) ({
  ^bb0(%block: tensor<i32>):
    %1 = "tfl.add"(%arg0, %arg0) { fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
    %2 = "tfl.add"(%1, %1) { fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
    "tfl.yield"(%2) : (tensor<i32>) -> ()
  },{

// CHECK: func @main([[MAIN_ARG0:%.*]]: tensor<i32>) -> tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK:      [[MAIN_TOKEN:%.*]] = mhlo.create_token

  // CHECK:      [[CALL:%.*]]:2 = call @callee([[MAIN_ARG0]], [[MAIN_TOKEN]])
  // CHECK-SAME: (tensor<i32>, !mhlo.token) -> (tensor<i32>, !mhlo.token)
  %0 = func.call @callee(%arg0) : (tensor<i32>) -> tensor<i32>

  // CHECK:      return [[CALL]]#0