// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: litert-opt %s -split-input-file -tfl-raise-custom-ops="test-raise-tf-targets=tf.FakeQuantWithMinMaxVarsPerChannel,tf.FakeQuantWithMinMaxVars" -tfl-prepare-tf | FileCheck --dump-input=always %s
// RUN: litert-opt %s -split-input-file -tfl-raise-custom-ops="test-raise-tf-targets=tf.FakeQuantWithMinMaxVarsPerChannel,tf.FakeQuantWithMinMaxVars" -tfl-prepare-tf=use-fake-quant-num-bits=true | FileCheck --check-prefix LOBIT --dump-input=always %s

// CHECK-LABEL: fakeQuantPerChannelForActivation
func.func @fakeQuantPerChannelForActivation(%arg0: tensor<8x4xf32>) -> (tensor<8x4xf32>) {
  %arg1 = arith.constant dense<[0.0, -1.0, 1.0, 0.0]> : tensor<4xf32>
  %arg2 = arith.constant dense<[255.0, 254.0, 256.0, 1.0e-9]> : tensor<4xf32>
  %0 = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<8x4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<8x4xf32>
  func.return %0 : tensor<8x4xf32>

// CHECK:  %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %cst, %cst_0)
// CHECK:  %[[q:.*]] = "tfl.quantize"(%[[fq]]) <{qtype = tensor<8x4x!quant.uniform<u8:f32:1, {1.000000e+00,1.000000e+00:1,1.000000e+00,3.9215686274509805E-9:127}>>}>
// CHECK:  %[[dq:.*]] = "tfl.dequantize"(%[[q]])
// CHECK:  return %[[dq]]
}

// -----

// CHECK-LABEL: fakeQuantForActivation
func.func @fakeQuantForActivation(tensor<8xf32>) -> (tensor<8xf32>) {
^bb0(%arg0: tensor<8xf32>):
  %arg1 = arith.constant dense<0.0> : tensor<f32>
  %arg2 = arith.constant dense<255.0> : tensor<f32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %0 : tensor<8xf32>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0)
// CHECK:  %1 = "tfl.quantize"(%0) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK:  %2 = "tfl.dequantize"(%1)
// CHECK:  return %2
}

// -----

// CHECK-LABEL: fakeQuantForActivationNoDuplication
func.func @fakeQuantForActivationNoDuplication(tensor<8xf32>) -> (tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>) {
^bb0(%arg0: tensor<8xf32>):
  %arg1 = arith.constant dense<0.0> : tensor<f32>
  %arg2 = arith.constant dense<255.0> : tensor<f32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>} : (tensor<8xf32>) -> tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>
  func.return %1 : tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 5 : i64}>
// CHECK:  %1 = "tfl.quantize"(%0) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK:  return %1
}

// -----

// CHECK-LABEL: WrappedFakeQuantFolded
func.func @WrappedFakeQuantFolded() -> tensor<8xf32> {
  %in = arith.constant dense<0.0> : tensor<8xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %mini = "tf.Identity"(%min) : (tensor<f32>) -> tensor<f32>
  %maxi = "tf.Identity"(%max) : (tensor<f32>) -> tensor<f32>
  %rst = "tfl.custom_tf"(%in, %mini, %maxi) ({
  ^bb0(%arg1: tensor<8xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>):
    %2 = "tf.FakeQuantWithMinMaxVars"(%arg1, %arg2, %arg3) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
    "tfl.yield"(%2) : (tensor<8xf32>) -> ()
  }) {num_bits = 5, narrow_range = false} :  (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %rst : tensor<8xf32>

// CHECK: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<8xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT]]) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: return %[[DEQUANTIZE]] : tensor<8xf32>
}

// -----

// CHECK-LABEL: fakeQuantFolded
func.func @fakeQuantFolded() -> (tensor<8xf32>) {
  %in = arith.constant dense<0.0> : tensor<8xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %mini = "tf.Identity"(%min) : (tensor<f32>) -> tensor<f32>
  %maxi = "tf.Identity"(%max) : (tensor<f32>) -> tensor<f32>
  %rst = "tf.FakeQuantWithMinMaxVars"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %rst : tensor<8xf32>

// CHECK: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<8xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT]]) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: return %[[DEQUANTIZE]] : tensor<8xf32>
}

// -----

// CHECK-LABEL: fakeQuantFoldedWithoutIdentity
func.func @fakeQuantFoldedWithoutIdentity() -> (tensor<8xf32>) {
  %in = arith.constant dense<0.0> : tensor<8xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %rst = "tf.FakeQuantWithMinMaxVars"(%in, %min, %max) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %rst : tensor<8xf32>

// CHECK: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<8xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT]]) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: return %[[DEQUANTIZE]] : tensor<8xf32>
}

// -----

// CHECK-LABEL: fakeQuantFoldedWithCast
func.func @fakeQuantFoldedWithCast() -> (tensor<8xf32>) {
  %in = arith.constant dense<0.0> : tensor<8xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %mini = "tf.Identity"(%min) : (tensor<f32>) -> tensor<f32>
  %maxi = "tf.Identity"(%max) : (tensor<f32>) -> tensor<f32>
  %minc = "tf.Cast"(%mini) : (tensor<f32>) -> tensor<f32>
  %maxc = "tf.Cast"(%maxi) : (tensor<f32>) -> tensor<f32>
  %rst = "tf.FakeQuantWithMinMaxVars"(%in, %minc, %maxc) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %rst : tensor<8xf32>

// CHECK: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<8xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT]]) <{qtype = tensor<8x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: return %[[DEQUANTIZE]] : tensor<8xf32>
}

// -----

// CHECK-LABEL: fakeQuantNotFolded
func.func @fakeQuantNotFolded(tensor<8xf32>, tensor<f32>, tensor<f32>) -> (tensor<8xf32>) {
^bb0(%arg0: tensor<8xf32>, %arg3: tensor<f32>, %arg4: tensor<f32>):
  %1 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg3, %arg4) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %1 : tensor<8xf32>

// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2)
// CHECK: return %0 : tensor<8xf32>
}

// -----

// CHECK-LABEL: fakeQuantFollowedByTranspose
func.func @fakeQuantFollowedByTranspose(tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> (tensor<2x1xf32>) {
^bb0(%arg0: tensor<1x2xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>):
  %cst_0 = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Transpose"(%0, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %1 : tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %0 = "tf.Transpose"(%arg0, %cst)
// CHECK:  %1 = "tf.FakeQuantWithMinMaxVars"(%0, %arg1, %arg2)
// CHECK:  return %1
}

// -----

// CHECK-LABEL: fakeQuantFollowedByTransposes
func.func @fakeQuantFollowedByTransposes(tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> (tensor<2x1xf32>, tensor<2x1xf32>) {
^bb0(%arg0: tensor<1x2xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>):
  %cst_0 = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Transpose"(%0, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  %2 = "tf.Transpose"(%0, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %[[FQ:.*]] = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2)
// CHECK:  %[[T1:.*]] = "tf.Transpose"(%[[FQ]], %cst)
// CHECK:  %[[T2:.*]] = "tf.Transpose"(%[[FQ]], %cst)
}

// -----

// CHECK-LABEL: fakeQuantFollowedByReshape
func.func @fakeQuantFollowedByReshape(tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> (tensor<2x1xf32>) {
^bb0(%arg0: tensor<1x2xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>):
  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1 : tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %0 = "tf.Reshape"(%arg0, %cst)
// CHECK-SAME: tensor<2x1xf32>
// CHECK:  %1 = "tf.FakeQuantWithMinMaxVars"(%0, %arg1, %arg2)
// CHECK:  return %1
}

// -----

// CHECK-LABEL: fakeQuantFollowedByReshapes
func.func @fakeQuantFollowedByReshapes(tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> (tensor<2x1xf32>, tensor<2x1xf32>) {
^bb0(%arg0: tensor<1x2xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>):
  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %[[FQ:.*]] = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2)
// CHECK:  %[[R1:.*]] = "tf.Reshape"(%[[FQ]], %cst)
// CHECK-SAME: tensor<2x1xf32>
// CHECK:  %[[R2:.*]] = "tf.Reshape"(%[[FQ]], %cst)
// CHECK-SAME: tensor<2x1xf32>
}

// -----

// CHECK-LABEL: fakeQuantWithConv2D
func.func @fakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x8x7x16xf32>) {
^bb0(%arg: tensor<256x32x32x3xf32>) :
  %in = arith.constant dense<0.0> : tensor<3x3x3x16xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %mini = "tf.Identity"(%min) : (tensor<f32>) -> tensor<f32>
  %maxi = "tf.Identity"(%max) : (tensor<f32>) -> tensor<f32>
  %fq = "tf.FakeQuantWithMinMaxVars"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<f32>, tensor<f32>) -> tensor<3x3x3x16xf32>
  %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32>
  func.return %rst : tensor<256x8x7x16xf32>

// CHECK-DAG: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<16xf32>}>
// CHECK-DAG: %[[CONSTANT0:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<16x3x3x3xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<16x3x3x3x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: %[[CONV:.*]] = "tfl.conv_2d"(%arg0, %[[DEQUANTIZE]], %[[CONSTANT]])
// CHECK: return %[[CONV]]
}

// -----

// CHECK-LABEL: perChannelFakeQuantWithConv2D
func.func @perChannelFakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x8x7x16xf32>) {
^bb0(%arg: tensor<256x32x32x3xf32>) :
  %in = arith.constant dense<0.0> : tensor<3x3x3x16xf32>
  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<255.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>
  %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32>
  func.return %rst : tensor<256x8x7x16xf32>

// CHECK-DAG: %[[CONSTANT0:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<16x3x3x3xf32>}>
// CHECK-DAG: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<16xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<16x3x3x3x!quant.uniform<u8:f32:0,
// CHECK-SAME: {1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00}>>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: %[[CONV:.*]] = "tfl.conv_2d"(%arg0, %[[DEQUANTIZE]], %[[CONSTANT]])
// CHECK: return %[[CONV]] : tensor<256x8x7x16xf32>
}

// -----

// CHECK-LABEL: fakeQuantWithDepthwiseConv2D
func.func @fakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x30x30x16xf32>) {
^bb0(%arg: tensor<256x32x32x3xf32>) :
  %in = arith.constant dense<0.0> : tensor<3x3x3x16xf32>
  %min = arith.constant dense<0.0> : tensor<f32>
  %max = arith.constant dense<255.0> : tensor<f32>
  %mini = "tf.Identity"(%min) : (tensor<f32>) -> tensor<f32>
  %maxi = "tf.Identity"(%max) : (tensor<f32>) -> tensor<f32>
  %fq = "tf.FakeQuantWithMinMaxVars"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<f32>, tensor<f32>) -> tensor<3x3x3x16xf32>
  %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32>
  func.return %rst : tensor<256x30x30x16xf32>

// CHECK-DAG: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<48xf32>}>
// CHECK-DAG: %[[CONSTANT0:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<1x3x3x48xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<1x3x3x48x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: %[[CONV:.*]] = "tfl.depthwise_conv_2d"(%arg0, %[[DEQUANTIZE]], %[[CONSTANT]])
// CHECK: return %[[CONV]]
}

// -----

// CHECK-LABEL: perChannelFakeQuantWithDepthwiseConv2D
func.func @perChannelFakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x30x30x16xf32>) {
^bb0(%arg: tensor<256x32x32x3xf32>) :
  %in = arith.constant dense<0.0> : tensor<3x3x3x16xf32>
  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<255.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>
  %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32>
  func.return %rst : tensor<256x30x30x16xf32>

// CHECK-DAG: %[[CONSTANT0:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<1x3x3x48xf32>}>
// CHECK-DAG: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<48xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<1x3x3x48x!quant.uniform<u8:f32:3,
// CHECK-SAME: {1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,
// CHECK-SAME:  1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,
// CHECK-SAME:  1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00}>>}
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: %[[CONV:.*]] = "tfl.depthwise_conv_2d"(%arg0, %[[DEQUANTIZE]], %[[CONSTANT]])
// CHECK: return %[[CONV]]
}

// -----

// CHECK-LABEL: perChannelFakeQuantWithDepthwiseConv2DWithReshape
func.func @perChannelFakeQuantWithDepthwiseConv2DWithReshape(%arg: tensor<1x160x160x48xf32>) -> (tensor<1x160x160x48xf32>) {
  %in = arith.constant dense<0.0> : tensor<3x3x48x1xf32>
  %min = arith.constant dense<0.0> : tensor<48xf32>
  %max = arith.constant dense<255.0> : tensor<48xf32>
  %mini = "tf.Identity"(%min) : (tensor<48xf32>) -> tensor<48xf32>
  %maxi = "tf.Identity"(%max) : (tensor<48xf32>) -> tensor<48xf32>
  %s1 = arith.constant dense<[3, 3, 48]> : tensor<3xi32>
  %s2 = arith.constant dense<[3, 3, 48, 1]> : tensor<4xi32>
  %r1 = "tf.Reshape"(%in, %s1) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48x1xf32>, tensor<3xi32>) -> tensor<3x3x48xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%r1, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<3x3x48xf32>
  %r2 = "tf.Reshape"(%fq, %s2) {T = f32, Tshape = i32, device = ""} : (tensor<3x3x48xf32>, tensor<4xi32>) -> tensor<3x3x48x1xf32>
  %rst = "tf.DepthwiseConv2dNative"(%arg, %r2) {T = f32, data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x160x160x48xf32>, tensor<3x3x48x1xf32>) -> tensor<1x160x160x48xf32>
  func.return %rst : tensor<1x160x160x48xf32>

// CHECK-DAG: %[[CONSTANT0:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<1x3x3x48xf32>}>
// CHECK-DAG: %[[CONSTANT:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<48xf32>}>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<1x3x3x48x!quant.uniform<u8:f32:3,
// CHECK-SAME: {1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,
// CHECK-SAME:  1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,
// CHECK-SAME:  1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00,1.000000e+00}>>}
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])
// CHECK: %[[CONV:.*]] = "tfl.depthwise_conv_2d"(%arg0, %[[DEQUANTIZE]], %[[CONSTANT]])
// CHECK: return %[[CONV]]
}

// -----

// LOBIT-LABEL: fakeQuant3BitPerChannelForActivation
func.func @fakeQuant3BitPerChannelForActivation(%arg0: tensor<8x4xf32>) -> (tensor<8x4xf32>) {
  %arg1 = arith.constant dense<[0.0, -1.0, -31.0, -30.0]> : tensor<4xf32>
  %arg2 = arith.constant dense<[31.0, 30.0, 31.0, 32.0]> : tensor<4xf32>
  %0 = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<8x4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<8x4xf32>
  func.return %0 : tensor<8x4xf32>

// LOBIT:  %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %cst, %cst_0)
// LOBIT:  %[[q:.*]] = "tfl.quantize"(%[[fq]]) <{qtype = tensor<8x4x!quant.uniform<u8<0:31>:f32:1, {1.000000e+00,1.000000e+00:1,2.000000e+00:16,2.000000e+00:15}>>}>
// LOBIT:  %[[dq:.*]] = "tfl.dequantize"(%[[q]])
// LOBIT:  return %[[dq]]
}

// -----

// LOBIT-LABEL: fakeQuant3BitForActivation
func.func @fakeQuant3BitForActivation(tensor<8xf32>) -> (tensor<8xf32>) {
^bb0(%arg0: tensor<8xf32>):
  %arg1 = arith.constant dense<-30.0> : tensor<f32>
  %arg2 = arith.constant dense<32.0> : tensor<f32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor<f32>, tensor<f32>) -> tensor<8xf32>
  func.return %0 : tensor<8xf32>

// LOBIT:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0)
// LOBIT:  %1 = "tfl.quantize"(%0) <{qtype = tensor<8x!quant.uniform<u8<0:31>:f32, 2.000000e+00:15>>}>
// LOBIT:  %2 = "tfl.dequantize"(%1)
// LOBIT:  return %2
}

// -----

// CHECK-LABEL: fakeQuantConcat
func.func @fakeQuantConcat(%arg0: tensor<1x6400x2xf32>, %arg1: tensor<1x1600x2xf32>) -> (tensor<1x8000x2xf32>) {
  %cst = arith.constant dense<1> : tensor<i32>
  %cst_1 = arith.constant dense<-1.0> : tensor<f32>
  %cst_2 = arith.constant dense<1.0> : tensor<f32>
  %cst_3 = arith.constant dense<-2.0> : tensor<f32>
  %cst_4 = arith.constant dense<0.5> : tensor<f32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_1, %cst_2) {num_bits = 8, narrow_range = false} : (tensor<1x6400x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x6400x2xf32>
  %1 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_3, %cst_4) {num_bits = 8, narrow_range = false} : (tensor<1x1600x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x1600x2xf32>
  %2 = "tf.ConcatV2"(%0, %1, %cst) : (tensor<1x6400x2xf32>, tensor<1x1600x2xf32>, tensor<i32>) -> tensor<1x8000x2xf32>
  return %2 : tensor<1x8000x2xf32>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_0, %cst_1)
// CHECK:  %1 = "tfl.quantize"(%0)
// CHECK:  %2 = "tfl.dequantize"(%1)
// CHECK:  %3 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_2, %cst_3)
// CHECK:  %4 = "tfl.quantize"(%3)
// CHECK:  %5 = "tfl.dequantize"(%4)
// CHECK:  %6 = "tf.ConcatV2"(%2, %5, %cst)
// CHECK:  %7 = "tf.FakeQuantWithMinMaxVars"(%6, %cst_2, %cst_1) <{narrow_range = false, num_bits = 8 : i64}> : (tensor<1x8000x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x8000x2xf32>
// CHECK:  %8 = "tfl.quantize"(%7) <{qtype = tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>}> : (tensor<1x8000x2xf32>) -> tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>
// CHECK:  %9 = "tfl.dequantize"(%8) : (tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>) -> tensor<1x8000x2xf32>
// CHECK:  return %9
}

// -----

// CHECK-LABEL: fakeQuantConcatQDQ
func.func @fakeQuantConcatQDQ(%arg0: tensor<1x6400x2xf32>, %arg1: tensor<1x1600x2xf32>) -> (tensor<1x8000x2xf32>) {
  %cst = arith.constant dense<1> : tensor<i32>
  %cst_1 = arith.constant dense<-1.0> : tensor<f32>
  %cst_2 = arith.constant dense<1.0> : tensor<f32>
  %cst_3 = arith.constant dense<-2.0> : tensor<f32>
  %cst_4 = arith.constant dense<0.5> : tensor<f32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_1, %cst_2) {num_bits = 8, narrow_range = false} : (tensor<1x6400x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x6400x2xf32>
  %1 = "tfl.quantize"(%0) {qtype = tensor<1x6400x2x!quant.uniform<u8:f32, 1.0>>} : (tensor<1x6400x2xf32>) -> tensor<1x6400x2x!quant.uniform<u8:f32, 1.0>>
  %2 = "tfl.dequantize"(%1) : (tensor<1x6400x2x!quant.uniform<u8:f32, 1.0>>) -> tensor<1x6400x2xf32>
  %3 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_3, %cst_4) {num_bits = 8, narrow_range = false} : (tensor<1x1600x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x1600x2xf32>
  %4 = "tfl.quantize"(%3) {qtype = tensor<1x1600x2x!quant.uniform<u8:f32, 1.0>>} : (tensor<1x1600x2xf32>) -> tensor<1x1600x2x!quant.uniform<u8:f32, 1.0>>
  %5 = "tfl.dequantize"(%4) : (tensor<1x1600x2x!quant.uniform<u8:f32, 1.0>>) -> tensor<1x1600x2xf32>
  %6 = "tf.ConcatV2"(%2, %5, %cst) : (tensor<1x6400x2xf32>, tensor<1x1600x2xf32>, tensor<i32>) -> tensor<1x8000x2xf32>
  return %6 : tensor<1x8000x2xf32>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_0, %cst_1)
// CHECK:  %1 = "tfl.quantize"(%0)
// CHECK:  %2 = "tfl.dequantize"(%1)
// CHECK:  %3 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_2, %cst_3)
// CHECK:  %4 = "tfl.quantize"(%3)
// CHECK:  %5 = "tfl.dequantize"(%4)
// CHECK:  %6 = "tf.ConcatV2"(%2, %5, %cst)
// CHECK:  %7 = "tf.FakeQuantWithMinMaxVars"(%6, %cst_2, %cst_1) <{narrow_range = false, num_bits = 8 : i64}> : (tensor<1x8000x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x8000x2xf32>
// CHECK:  %8 = "tfl.quantize"(%7) <{qtype = tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>}> : (tensor<1x8000x2xf32>) -> tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>
// CHECK:  %9 = "tfl.dequantize"(%8) : (tensor<1x8000x2x!quant.uniform<u8:f32, 0.011764705882352941:170>>) -> tensor<1x8000x2xf32>
// CHECK:  return %9
}

// -----

// CHECK-LABEL: populateFakeQuantOnMeanOutput
func.func @populateFakeQuantOnMeanOutput(%arg0: tensor<f32>) -> (tensor<f32>) {
  %cst = arith.constant dense<-1.0> : tensor<f32>
  %cst_1 = arith.constant dense<1.0> : tensor<f32>
  %cst_2 = arith.constant dense<0> : tensor<1xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_1) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %1 = "tf.Mean"(%0, %cst_2) <{keep_dims = false}> : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
  return %1 : tensor<f32>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0)
// CHECK:  %1 = "tfl.quantize"(%0) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK:  %2 = "tfl.dequantize"(%1) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK:  %3 = "tf.Mean"(%2, %cst_1)
// CHECK:  %4 = "tf.FakeQuantWithMinMaxVars"(%3, %cst, %cst_0)
// CHECK:  %5 = "tfl.quantize"(%4) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK:  %6 = "tfl.dequantize"(%5) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK:  return %6
}

// -----

// CHECK-LABEL: populateFakeQuantOnMeanOutputQDQs
func.func @populateFakeQuantOnMeanOutputQDQs(%arg0: tensor<f32>) -> (tensor<f32>) {
  %cst = arith.constant dense<-1.0> : tensor<f32>
  %cst_1 = arith.constant dense<1.0> : tensor<f32>
  %cst_2 = arith.constant dense<0> : tensor<1xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_1) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %1 = "tfl.quantize"(%0) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
  %2 = "tfl.dequantize"(%1) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
  %3 = "tf.Mean"(%2, %cst_2) <{keep_dims = false}> : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
  return %3 : tensor<f32>

// CHECK:       %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0)
// CHECK-NEXT:  %1 = "tfl.quantize"(%0) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %2 = "tfl.dequantize"(%1) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  %3 = "tf.Mean"(%2, %cst_1)
// CHECK-NEXT:  %4 = "tf.FakeQuantWithMinMaxVars"(%3, %cst, %cst_0)
// CHECK-NEXT:  %5 = "tfl.quantize"(%4) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %6 = "tfl.dequantize"(%5) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  return %6
}

// -----

// CHECK-LABEL: populateFakeQuantOnMeanOutputFollowedByConcat
func.func @populateFakeQuantOnMeanOutputFollowedByConcat(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<1xf32>) {
  %cst = arith.constant dense<1> : tensor<i32>
  %cst_1 = arith.constant dense<-1.0> : tensor<f32>
  %cst_2 = arith.constant dense<1.0> : tensor<f32>
  %cst_3 = arith.constant dense<0> : tensor<1xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_1, %cst_2) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %1 = "tf.Mean"(%0, %cst_3) <{keep_dims = false}> : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
  %2 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_1, %cst_2) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = "tf.Mean"(%2, %cst_3) <{keep_dims = false}> : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
  %4 = "tf.ConcatV2"(%1, %3, %cst) : (tensor<f32>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
  return %4 : tensor<1xf32>

// CHECK:       %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst_0, %cst_1)
// CHECK-NEXT:  %1 = "tfl.quantize"(%0) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %2 = "tfl.dequantize"(%1) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  %3 = "tf.Mean"(%2, %cst_2)
// CHECK-NEXT:  %4 = "tf.FakeQuantWithMinMaxVars"(%3, %cst_0, %cst_1)
// CHECK-NEXT:  %5 = "tfl.quantize"(%4) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %6 = "tfl.dequantize"(%5) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  %7 = "tf.FakeQuantWithMinMaxVars"(%arg1, %cst_0, %cst_1)
// CHECK-NEXT:  %8 = "tfl.quantize"(%7) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %9 = "tfl.dequantize"(%8) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  %10 = "tf.Mean"(%9, %cst_2)
// CHECK-NEXT:  %11 = "tf.FakeQuantWithMinMaxVars"(%10, %cst_0, %cst_1)
// CHECK-NEXT:  %12 = "tfl.quantize"(%11) <{qtype = tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<f32>) -> tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %13 = "tfl.dequantize"(%12) : (tensor<!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<f32>
// CHECK-NEXT:  %14 = "tf.ConcatV2"(%6, %13, %cst)
// CHECK-NEXT:  %15 = "tf.FakeQuantWithMinMaxVars"(%14, %cst_0, %cst_1) <{narrow_range = false, num_bits = 8 : i64}> : (tensor<1xf32>, tensor<f32>, tensor<f32>) -> tensor<1xf32>
// CHECK-NEXT:  %16 = "tfl.quantize"(%15) <{qtype = tensor<1x!quant.uniform<u8:f32, 0.0078431372549019607:128>>}> : (tensor<1xf32>) -> tensor<1x!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK-NEXT:  %17 = "tfl.dequantize"(%16) : (tensor<1x!quant.uniform<u8:f32, 0.0078431372549019607:128>>) -> tensor<1xf32>
// CHECK-NEXT:  return %17
}

// -----

// CHECK-LABEL: populateFakeQuantOnMeanOutputNegativeCase
func.func @populateFakeQuantOnMeanOutputNegativeCase(%arg0: tensor<f32>) -> (tensor<f32>) {
  %cst = arith.constant dense<-1.0> : tensor<f32>
  %cst_1 = arith.constant dense<1.0> : tensor<f32>
  %cst_2 = arith.constant dense<0> : tensor<1xi32>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_1) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %1 = "tf.Mean"(%0, %cst_2) <{keep_dims = false}> : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
  %2 = "tf.FakeQuantWithMinMaxVars"(%1, %cst, %cst_1) {num_bits = 8, narrow_range = false} : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  return %2 : tensor<f32>

// CHECK:      %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0)
// CHECK:      %1 = "tfl.quantize"(%0)
// CHECK:      %2 = "tfl.dequantize"(%1)
// CHECK:      %3 = "tf.Mean"(%2, %cst_1)
// CHECK:      %4 = "tf.FakeQuantWithMinMaxVars"(%3, %cst, %cst_0)
// CHECK-NOT:  "tf.FakeQuantWithMinMaxVars"
}

// -----