auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
  if (!output_type) return failure();

  auto input_quantized_type =
      quant::QuantizedType::getQuantizedElementType(input_type);
  auto output_quantized_type =
      quant::QuantizedType::getQuantizedElementType(output_type);
  // If both the input & output types are non-quantized, they will be both
  // nullptrs.

  // Matches all convolution quantizable unit family.
  spec.mutable_matcher()->mutable_function_name()->set_regex(
      "composite_conv.*");

  // Enable per-channel quantization for convolution weights.
  QuantizedType conv_weight_quantized_type{};

  // Assumes NHWC format, specifying the channel dimension (3) as the
  // quantized axis.
  conv_weight_quantized_type.mutable_dimension_specs()->set_dimension(3);

// index.
template <typename LstmOp>
inline QuantizedType GetIntermediateElementType(LstmOp op, int tensor_index) {
  if (tensor_index < 0 || tensor_index > 4) return nullptr;
  TypeAttr attr = op->template getAttrOfType<TypeAttr>(
      intermediate_attributes[tensor_index]);
  if (!attr) {
    return nullptr;
  }
  return QuantizedType::getQuantizedElementType(attr.getValue());
}

  if (mlir::isa<FloatType>(type.getElementType())) {
    for (const auto val : attr.getValues<APFloat>()) {
      if (val.isZero()) {
        num_zeros++;
      }
    }
  } else if (mlir::isa<quant::QuantizedType>(type.getElementType())) {
    for (const auto val : attr.getValues<int8_t>()) {
      if (val == 0) {
        num_zeros++;
      }
    }
  }

  return GetSparsity(num_zeros, num_elements);
}

    if (!type.isInteger(32) && !type.isInteger(64)) return false;
  } else {
    // QConstOp path.
    auto qconst_op = dyn_cast<TFL::QConstOp>(op);
    auto type =
        quant::QuantizedType::getQuantizedElementType(qconst_op.getType());
    if (type.getStorageTypeIntegralWidth() != 32) {
      return false;
    }
  }
  return true;
}

void FoldConstantsToSubgraphPass::runOnOperation() {

}
// CHECK-SAME: %[[ARG_0:.+]]: tensor<1x3x2x3xf32>

// Test that the weight is prepared for per-tensor quantization, based on the
// `_quantization_method` attribute without a `dimension_specs` field in
// `QuantizedType`.
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} tensor<2x3x3x2xf32>

// at operation runtime.
//
// <type> can be:
//   "string", "int", "float", "bool", "type", "shape", or "tensor"
//   "numbertype", "realnumbertype", "quantizedtype"
//       (meaning "type" with a restriction on valid values)
//   "{int32,int64}" or {realnumbertype,quantizedtype,string}"
//       (meaning "type" with a restriction containing unions of value types)
//   "{\"foo\", \"bar\n baz\"}", or "{'foo', 'bar\n baz'}"

    return success();
  }
};

quant::QuantizedType getQuantizedElementType(CastOp cast_op) {
  if (!cast_op || !cast_op.getInputElementType()) {
    return {};
  }
  return cast_op.getInputElementType()
      .cast<TypeAttr>()
      .getValue()
      .dyn_cast<quant::QuantizedType>();
}

class RemoveRawDataOp : public OpRewritePattern<TFRQuantRawDataOp> {

  // operand of `quantize_op`.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      QuantizeOpT quantize_op) const {
    Value operand = quantize_op->getOperand(0);
    if (QuantizedType::getQuantizedElementType(operand.getType())) {
      // The input of the quantize op has already been quantized, i.e.
      // rescale.
      return failure();
    }

#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"

namespace mlir {
namespace quant {

static bool EmptyParams(QuantParams p) { return p == quant::QuantizedType(); }

// The state for each op result during the quantization parameters propagation.
struct QuantState {
  // Quantization parameters propagated to an op result.
  QuantParams params;