func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
      func.return %arg0 : tensor<?xi32>
    }

    func.func @main(%arg0: tensor<20xi32>, %arg1: tensor<?xi32>) -> tensor<*xi32> {
      %2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>
      func.return %2 : tensor<*xi32>
    }
  })";

func.func private @simple_callee(%arg0: tensor<2xi32>) -> tensor<2xi32> {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<2xi32>
  }
  func.return %graph#0 : tensor<2xi32>
}

// Test with a TPUPartitionedCallOp that is not inlined.
func.func @dont_inline_tpu_partitioned_call(%arg0: tensor<2xi32> {tf._user_specified_name = "inputs_0"}, %arg1: tensor<2xi32> {tf._user_specified_name = "inputs_1"}) -> tensor<2xi32> {

>, tensor<i8>, tensor<f64>, tensor<!tf_type.resource>)

# CHECK-SUBTYPE: func @main(%arg0: tensor<10xi32>, %arg1: tensor<!tf_type.variant<tensor<10xf32>>>, %arg2: tensor<10xi8>, %arg3: tensor<10xf64>, %arg4: tensor<!tf_type.resource<tensor<10xi32>>>) -> (tensor<10xi32>, tensor<!tf_type.variant<tensor<10xf32>>>, tensor<10xi8>, tensor<10xf64>, tensor<!tf_type.resource<tensor<10xi32>>>)...

module {
  func.func @same_scale_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
    %cst_1 = arith.constant dense<1.0> : tensor<144x10xf32>
    %cst_2 = arith.constant dense<0.1> : tensor<10xf32>
    %0 = "quantfork.qcast"(%arg0) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.05:-10>>

  func.return %0#0 : tensor<2xi32>
}

// CHECK-LABEL: mirror_pad
func.func @mirror_pad(%arg0: tensor<2x3xcomplex<f64>>) -> tensor<4x7xcomplex<f64>> {
  %0 = mhlo.constant dense<[[1, 1], [2, 2]]> : tensor<2x2xi32>
  // NO_FALLBACK: tf.MirrorPad
  // SUPPORTED_FALLBACK_DEVICE-NOT: tf.MirrorPad

}

// -----

func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
  // CHECK: tac.cost = 1.000000e+03
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  func.return %0 : tensor<10x10x10xf32>
}

// -----

    %cst_0 = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi32>} : () -> tensor<2xi32>
    %1 = "tf.MatMul"(%arg0, %arg1) {
      transpose_a = false, transpose_b = false
    } : (tensor<1x10xf32>, tensor<10x10xf32>) -> tensor<1x10xf32>  loc(fused["MatMul:", "test_opt_out"])

  %m0 = mhlo.constant dense<0> : tensor<i32>
  %t0 = "tf.Const"() { value = dense<0> : tensor<1xi32> } : () -> tensor<1xi32>
  %m1 = mhlo.constant dense<1> : tensor<i32>
  %t1 = "tf.Const"() { value = dense<1> : tensor<1xi32> } : () -> tensor<1xi32>
  %m2 = mhlo.constant dense<1> : tensor<i32>
  %t2 = "tf.Const"() { value = dense<1> : tensor<1xi32> } : () -> tensor<1xi32>
  %m3 = mhlo.constant dense<1> : tensor<i32>

func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  %1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>

    %4 = "quantfork.stats"(%2) {layerStats = dense<[-2795.36523, 4609.57373]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32>
    %5 = "tf.BiasAdd"(%4, %3) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %6 = "quantfork.stats"(%5) {layerStats = dense<[-2788.31055, 4616.62842]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32>