%outputs_1, %outputs_2 = "tf.IdentityN"(%outputs, %outputs_0) {device = "TPU"} : (tensor<2x2xf32>, tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>)
  func.return %outputs_0, %outputs_1 : tensor<2x2xf32>, tensor<2x2xf32>
}

// CHECK-LABEL: func @keep_identity_n_test
func.func @keep_identity_n_test(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
    // CHECK: tf.MatMul

  %noncon3 = "tf.Square"(%noncon2) : (tensor<2xf32>) -> tensor<2xf32>

  %res1 = "tf.DivNoNan"(%arg0, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  %res2 = "tf.MulNoNan"(%arg0, %noncon2) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  %res3 = "tf.DivNoNan"(%arg0, %noncon3) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  func.return %res1, %res2, %res3 : tensor<2xf32>, tensor<2xf32>, tensor<2xf32>

    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_b", device = "", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
    return %0 : tensor<2x2xf32>
  }

}

// CHECK-LABEL: pack_multiple
func.func @pack_multiple(%arg0: tensor<2x3xf32>,
                    %arg1: tensor<2x3xf32>,
                    %arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[AXIS:.*]] = arith.constant 0 : i32

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

func.func @my_add_n_single_input(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
  %0 = "tf.MyAddN"(%arg0) {N=1:i32} : (tensor<2x3xf32>) -> tensor<2x3xf32>
  func.return %0 : tensor<2x3xf32>

// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}

// CHECK-LABEL: my_add_n_multiple_inputs
func.func @my_add_n_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<2x3xf32> {

    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
      %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
      %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>

       {
         data_format = "NCHW",
         epsilon = 1.001000e-05 : f32
       } : (tensor<?x256x56x56xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>)
        -> (tensor<?x256x56x56xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<*xf32>)

  // CHECK: %[[BATCH_NORM1:[_a-z0-9]*]], {{.*}} = "tf.FusedBatchNormV3"
  // CHECK-SAME: %[[CONV1]]

    } : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
    %2 = "quantfork.stats"(%1) {layerStats = dense<[0.000000e+00, 6.000000e+00]> : tensor<2xf32>} : (tensor<2x2xf32>) -> tensor<2x2xf32>
    return %2 : tensor<2x2xf32>
  }

  // CHECK-LABEL: composite_dot_general
  func.func private @composite_dot_general(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {

  %1 = "quantfork.stats"(%0) {bitsNum = 8 : i64, layerStats = dense<[-2.0, 2.0]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>
  %2 = stablehlo.convert %1 : (tensor<2x3xf32>) -> (tensor<2x3xf32>)
  func.return %2 : tensor<2x3xf32>
}

// -----

// CHECK-LABEL: func @dot_redundant_stats