return ExpressedToQuantizedConverter{input_type, element_type};
  }
  // Supported primitive type (which just is the expressed type).
  if (isQuantizablePrimitiveType(input_type))
    return ExpressedToQuantizedConverter{input_type, input_type};
  // Unsupported.
  return ExpressedToQuantizedConverter{input_type, nullptr};
}

Type ExpressedToQuantizedConverter::convert(

    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
      if (input_type == nullptr || !input_type.hasStaticShape()) continue;
      // Quantized type does not support getSizeInBits.
      if (IsQUI8Type(input_type) || IsQI8Type(input_type)) {
        total_size_transferred += input_type.getNumElements() * 8;
      } else {

    int64_t total_count = 0;
    for (auto input : op->getOperands()) {
      auto input_type =
          mlir::dyn_cast_or_null<mlir::RankedTensorType>(input.getType());
      if (!input_type || !input_type.hasStaticShape()) {
        return false;
      }
      total_count += input_type.getNumElements();
    }
    *count = total_count;
    return true;
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ModifyIONodesPass)

  explicit ModifyIONodesPass() {}
  explicit ModifyIONodesPass(mlir::Type input_type, mlir::Type output_type) {
    this->input_type = input_type;
    this->output_type = output_type;
  }

  void runOnOperation() override;

 private:
  // Assign the io types from the command line flag. This is only required for
  // tests.

    SmallVector<Value, 4> dequantized_inputs;
    for (auto& input : op->getOpOperands()) {
      auto input_type = input.get().getType();
      if (IsQI8Type(input_type) || IsQUI8Type(input_type) ||
          IsQI32Type(input_type)) {
        auto dequantized_input_type =
            mlir::quant::QuantizedType::castToExpressedType(input_type);
        builder->setInsertionPoint(op);
        auto dequantize_op = builder->create<TFL::DequantizeOp>(

    // can have smaller memory usage.
    auto input_type =
        mlir::dyn_cast<RankedTensorType>(dequantize_op.getOutput().getType());
    auto output_type = mlir::dyn_cast<RankedTensorType>(
        passthrough_op->getResult(0).getType());
    if (!input_type || !output_type ||
        get_num_elements(input_type) <= get_num_elements(output_type)) {
      return failure();
    }

  if (!dims_type) return success();
  if (dims_type.getRank() > 1)
    return emitError(loc, "dimensions can only be 0D or 1D tensor");

  auto input_type = mlir::dyn_cast<RankedTensorType>(input.getType());
  if (!input_type) return success();
  int64_t rank = input_type.getRank();

  DenseIntElementsAttr dims_attr;
  if (!matchPattern(dims, m_Constant(&dims_attr))) return success();

    // Create a tfl.transpose op that performs ZX transpose on `input`.
    auto create_z_x_transpose_op = [&](Value input) -> Value {
      RankedTensorType input_type =
          mlir::cast<RankedTensorType>(input.getType());
      const int input_rank = input_type.getRank();

      // Create a 1D I32 tensor for representing the dimension permutation.
      auto permuation_tensor_type =

  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
        xla_shape.dimensions(), &tensor_shape));
    arg_shapes.emplace_back(tensor_shape);

    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

    int quant_dim = -1;
    auto input_type = mlir::cast<ShapedType>(input.getType());
    if (PerAxis) {
      if (!input_type.hasRank()) {
        tf_op.emitError("The input should have known rank for per-channel op.");
        return failure();
      }
      // This is a special case that the quant_dim is the last dimensions.
      quant_dim = input_type.getRank() - 1;
    }