//              tf.dequantize
//                   |
template <typename TFFakeQuantOp, bool PerAxis>
struct InsertQuantOpsAfterTFFakeQuantOp
    : public OpRewritePattern<TFFakeQuantOp> {
  using BaseType = InsertQuantOpsAfterTFFakeQuantOp<TFFakeQuantOp, PerAxis>;

  explicit InsertQuantOpsAfterTFFakeQuantOp<TFFakeQuantOp, PerAxis>(
      MLIRContext *ctx)
      : OpRewritePattern<TFFakeQuantOp>(ctx) {}

// TF::FakeQuant operators
using PreparePerTensorFakeQuant = ConvertFakeQuantOpToQuantOps<
    TF::FakeQuantWithMinMaxVarsOp, /*PerAxis=*/false,
    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsOp>>;

using PreparePerChannelFakeQuant = ConvertFakeQuantOpToQuantOps<
    TF::FakeQuantWithMinMaxVarsPerChannelOp, /*PerAxis=*/true,
    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsPerChannelOp>>;

using PreparePerTensorFakeQuant = InsertTFLQuantOpsAfterTFFakeQuantOp<
    TF::FakeQuantWithMinMaxVarsOp, /*PerAxis=*/false,
    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsOp>>;

using PreparePerChannelFakeQuant = InsertTFLQuantOpsAfterTFFakeQuantOp<
    TF::FakeQuantWithMinMaxVarsPerChannelOp, /*PerAxis=*/true,
    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsPerChannelOp>>;

// case of back-to-back tf.FakeQuant occurs
//
//             tf.FakeQuant*
//                   |
//             tf.FakeQuant*
//
template <typename TFFakeQuantOp, bool PerAxis, class FetchMinMax>
class InsertTFLQuantOpsAfterTFFakeQuantOp {
 public:
  explicit InsertTFLQuantOpsAfterTFFakeQuantOp(bool use_fake_quant_num_bits)
      : use_fake_quant_num_bits_(use_fake_quant_num_bits) {}

// case of back-to-back tf.FakeQuant occurs
//
//             tf.FakeQuant*
//                   |
//             tf.FakeQuant*
//
template <typename TFFakeQuantOp, bool PerAxis, class FetchMinMax>
class ConvertFakeQuantOpToQuantOps {
 public:
  explicit ConvertFakeQuantOpToQuantOps(bool use_fake_quant_num_bits)
      : use_fake_quant_num_bits_(use_fake_quant_num_bits) {}

  }
  for (auto spec : rop.getOutputSpecs()) {
    out_specs.push_back(spec.cast<TypeAttr>().getValue());
  }

  auto in_spec = input_specs[0].dyn_cast<UniformQuantizedType>();
  // TODO(fengliuai): handles the PerAxis QuantizedType.
  auto w_spec = input_specs[1].dyn_cast<UniformQuantizedType>();
  auto b_spec = input_specs[2].dyn_cast<UniformQuantizedType>();
  auto o_spec = out_specs[0].dyn_cast<UniformQuantizedType>();

// The following functions tests example quantization patterns outputted from
// JAX Quantizer. JAX Quantizer should output integer types, which are
// composed into `UniformQuantized{|PerAxis}Type` via
// `compose_uniform_quantized_type_pass.cc`.
// ============================================================================