return success();
  }
};

// UniformDequantizeOp takes TF quantized types as input which would have been
// converted to the mhlo quantized types. Use OpConversionPattern in order to
// retrieve the operand type *after* conversion, using OpAdaptor operand
// accessor.
// Same for other Uniform Quant Ops that take TF quantized types as input.
class ConvertUniformDequantizeOp

        per_axis_type.getStorageTypeMin(), per_axis_type.getStorageTypeMax());
  }

  auto quantize = builder.create<quantfork::QuantizeCastOp>(
      q_op.getLoc(), new_value_type.clone(new_qtype), new_value);
  auto dequantize = builder.create<quantfork::DequantizeCastOp>(
      dq_op.getLoc(), new_value_type, quantize.getResult());
  return dequantize.getResult();
}

    "quantfork::QuantizationForkDialect"
  ];
  let options = [
      ListOption<"quantize_allowlist_", "quantize-allowlist", "std::string",
                 "comma separated list of allowlisted functions to be quantized. Only used in tests">,
      Option<"quantize_signed_", "quantize-signed", "bool", "false",
             "signed inference type. Only used in tests">,

class UsedBy<string op> : Constraint<
  CPred<"llvm::isa<mlir::TFL::" # op # "Op>(*$0.getUsers().begin())">>;

// When the op is passing-through, the output types of the quantized ops need
// to be updated as well. Since the quantize op manages its own type by the
// "qtype" attribute, we should update the type shape in this attribute.
def ReorderTransposeDequantQuant :
      Pat<(TF_TransposeOp:$old_value

constexpr StringRef kCompositeFuncPrefix = "composite_";
constexpr StringRef kEmptyNodeName = "_empty_node";

// Returns a pair: `func_name` and `node_name` for the lifted function. In TF
// quantizer, both are filled. For StableHLO quantizer, the func_name is only
// filled and node_name is always set to "_empty_node".
std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {

  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
  ASSERT_THAT(func_op, NotNull());

  auto uniform_quantize_op_itr =
      func_op.getBody().op_begin<mlir::stablehlo::UniformQuantizeOp>();
  ASSERT_THAT(
      uniform_quantize_op_itr,
      Ne(func_op.getBody().op_end<mlir::stablehlo::UniformQuantizeOp>()));

  // `uniform_quantize` is considered partially quantized because its output is

        return success();
      }

      op.replaceAllUsesWith(q.getInput());
      return success();
    }
    return failure();
  }
};

// Fold the constant quantized Transpose ops.
struct FoldTransposeOp : public OpRewritePattern<TransposeOp> {
  explicit FoldTransposeOp(MLIRContext* context)
      : OpRewritePattern<TransposeOp>(context, 1) {}

    return %1 : tensor<1x3xf32>
  }
}

// -----

// Do not merge when function is not quantized function.

module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: func.func private @merge_relu_fusion
  func.func private @merge_relu_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {

          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  // Initialize for tests.
  void initializeForTest() {
    if (!test_mode_) return;

    op_set_.setCallback([this](const OpSet& new_op_set) {
      quant_options_.set_op_set(new_op_set);

    }

    // TODO: b/264218457 - Refactor the component below once StableHLO Quantizer
    // can run DRQ. Temporarily using TF Quantization for StableHLO DRQ.
    if (!toco_flags.has_quantization_options()) {
      // The default minimum number of elements a weights array must have to be
      // quantized by this transformation.
      const int kWeightsMinNumElementsDefault = 1024;