func.return %0 : tensor<1xi32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32> {tf_saved_model.index_path = ["input"]}) -> (tensor<1xi32> {tf_saved_model.index_path = ["r"]})
  attributes {tf_saved_model.exported_names = ["main"]} {

    %4 = "tf.LayerNorm"(%a1, %a2, %a3, %a4) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>
   "tfl.yield"(%4) : (tensor<128x128xf32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>

    %4 = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
    %5 = "tf.Const"() {value = dense<2.500000e-01> : tensor<f32>} : () -> tensor<f32>
    %6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
    %7 = "tf.Const"() {value = dense<[-1, 1001]> : tensor<2xi32>} : () -> tensor<2xi32>
    %8 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>

  // CHECK-NEXT: tf_mlrt.executeop([[y]]
  %v = "tf.AddV2"(%y, %z) {__op_key = 3: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
  return %w, %u, %v : tensor<1xi32>, tensor<1xi32>, tensor<1xi32>
}

// -----

// Test node names are preserved.

// CHECK-LABEL: func @main
func.func @main(%x: tensor<i32>) -> tensor<i32> {
  // CHECK: tf_mlrt.executeop

    %1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>

    // Use identity op to avoid the weight being constant-folded.
    %identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
    %2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
    %3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>

    %5 = "tf.MatMul"(%1, %3) {

// CHECK-SAME: (tensor<1x2x2x3xi8>, tensor<2x2x3x2xi8>, tensor<2xi32>, tensor<f32>, tensor<i32>, tensor<2xf32>, tensor<2xi32>, tensor<2xf32>, tensor<2xi32>, tensor<f32>, tensor<i32>) -> tensor<*xf32>
// CHECK: return %[[conv_quant]]

// CHECK-LABEL: func private @quantized_conv2d_with_bias_and_relu6_float_output_fn_0

  return %0 : tensor<*xi32>
}
func.func private @_batched(%arg0: tensor<1xi32>) -> tensor<1xi32> {
  return %arg0 : tensor<1xi32>
}

}

// -----

module attributes {tf_saved_model.semantics} {

// Test hoisting write side-effect ops.

// CHECK-DAG:       %[[VAL_3:.*]] = "tfl.pseudo_const"{{.*}}dense<2> : tensor<1xi32>
// CHECK-DAG:       %[[VAL_4:.*]] = "tfl.pseudo_const"{{.*}}dense<[2, 1]> : tensor<2xi32>
// CHECK:           %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:           %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>

  %2 = "tfl.add"(%arg0, %arg3) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %3 = "tfl.pack"(%1, %2) {tac.device = "CPU", tac.inference_type = "FLOAT", axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  return %3 : tensor<2x1xf32>
}
```

module {
  func.func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
    %0 = func.call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
    %1 = func.call @func_1_GPU_FLOAT(%arg0, %arg3) {tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>