func.return %4 : tensor<*xf32>
}

// CHECK: %[[q:.*]] = "quantfork.qcast"(%arg0)
// CHECK: %[[sc1:.*]] = "quantfork.scast"(%[[q]]) : (tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>)
// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[sc1]]) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32>
// CHECK: %[[avgpool_f32:.*]] = "tf.AvgPool"(%[[fcast]])
// CHECK-SAME: (tensor<*xf32>) -> tensor<*xf32>
// CHECK: %[[round:.*]] = "tf.Round"(%[[avgpool_f32]])

  addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.cc.inc"
      >();
}

OpFoldResult StorageCastOp::fold(FoldAdaptor) {
  // Matches x -> [scast -> scast] -> y, replacing the second scast with the
  // value of x if the casts invert each other.
  auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
  if (!srcScastOp || srcScastOp.getArg().getType() != getType())

  addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc.inc"
      >();
}

OpFoldResult StorageCastOp::fold(FoldAdaptor) {
  // Matches x -> [scast -> scast] -> y, replacing the second scast with the
  // value of x if the casts invert each other.
  auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
  if (!srcScastOp || srcScastOp.getArg().getType() != getType())

// CHECK: %[[maxpool:.*]] = "tf.MaxPool"
// CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[maxpool]])
// CHECK-SAME: quant.uniform<i8:f32, 5.000000e-02:-10>
// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])
// CHECK-SAME: quant.uniform<i8:f32, 5.000000e-02:-10>
// CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[dq1]]
// CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[reshape]])
// CHECK-SAME: quant.uniform<i8:f32, 5.000000e-02:-10>

// CHECK: %[[q0:.*]] = "quantfork.qcast"(%arg0)
// CHECK: %[[dq0:.*]] = "quantfork.dcast"(%[[q0]])
// CHECK-SAME: quant.uniform<i8:f32, 0.010039215461880554:-1>
// CHECK: %[[maxpool:.*]] = "tf.MaxPool"(%[[dq0]])
// CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[maxpool]])
// CHECK-SAME: quant.uniform<i8:f32, 0.010039215461880554:-1>
// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])

}

// CHECK: %[[q1:.*]] = "quantfork.qcast"(%arg0)
// CHECK-SAME: tensor<3x!quant.uniform<i8:f32, 0.0076314610593459188:-3>>
// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])
// CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg1)
// CHECK-SAME: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {0.003937007874015748,0.0039370079913477263:-25,0.003937007874015748:51}>>
// CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])

// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])

// CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg0)
// CHECK-SAME: tensor<1x3x4x3x!quant.uniform<i8:f32, 0.58754816990272674:-128>>
// CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])

// CHECK: %[[call:.*]] = "tf.PartitionedCall"(%[[dq2]], %[[dq1]], %[[dq0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_10
// CHECK: %[[q3:.*]] = "quantfork.qcast"(%[[call]]) {volatile}

    // CHECK: %[[q_weight:.*]] = "quantfork.qcast"
    // CHECK-SAME: -> tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {0.049663885371891529,0.060200210631363035}>>
    // CHECK: %[[dq_weight:.*]] = "quantfork.dcast"(%[[q_weight]])
    %cst = "tf.Const"() {device = "", value = dense<[[-6.30731344, 5.4962182], [1.80364347, -7.64542675]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>
    // CHECK: %[[q_act:.*]] = "quantfork.qcast"(%arg0)

// CHECK:  %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %cst, %cst_0)
// CHECK:  %[[q:.*]] = "quantfork.qcast"(%[[fq]]) : (tensor<8x3xf32>) -> tensor<8x3x!quant.uniform<i4:f32:1, {1.000000e+00:-8,1.000000e+00:-7,1.000000e+00:-8}>>
// CHECK:  %[[dq:.*]] = "quantfork.dcast"(%[[q]])
// CHECK:  return %[[dq]]
}

// CHECK-LABEL: fakeQuantForActivation

// CHECK:  %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg0, %cst, %cst_0)
// CHECK:  %[[q:.*]] = "quantfork.qcast"(%[[fq]]) : (tensor<8x3xf32>) -> tensor<8x3x!quant.uniform<i8:f32:1, {1.000000e+00:-128,1.000000e+00:-127,1.000000e+00:-128}>>
// CHECK:  %[[dq:.*]] = "quantfork.dcast"(%[[q]])
// CHECK:  return %[[dq]]
}

// CHECK-LABEL: fakeQuantForActivation