inline constexpr StringRef kNullAttributeValue = "N/A";

// Prefixes attached to lifted functions.
constexpr StringRef kQuantizedFuncPrefix = "quantized_";
constexpr StringRef kCompositeFuncPrefix = "composite_";

// The attribute will be used for TF::XlaCallModuleOp to restore the original
// function name when loading it back.
inline constexpr StringRef kOriginalStablehloEntryFunctionAttrName =
    "_original_entry_function";

  });
}

// Populates non-quantized ops from `module_op` to `results`. After going
// through the quantization passes, non-quantized quantizable units remain as
// `TF::XlaCallModuleOp` with a callee's prefix of `composite_`.
void PopulateNonQuantizedResults(ModuleOp module_op,
                                 QuantizationResults& results) {
  module_op.walk([&results](TF::XlaCallModuleOp xla_call_module_op) {

  if (auto call_op = dyn_cast<TF::PartitionedCallOp>(op)) {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }
    if (function_name.contains("depthwise_conv2d")) {
      spec->coeff_op_quant_dim[1] = 3;
      if (function_name.contains("with_bias")) {
        spec->biases_params[2] = {{0, 1},

    auto entry_function =
        call_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
    StringRef function_name = entry_function.getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }

    if (function_name.contains("conv")) {
      // Looks up `Method` to see if it should be per-channel quantized and
      // populates the spec accordingly.

using DebuggerType = DebuggerConfig::DebuggerType;

constexpr StringRef kEntryFuncAttrName = "_entry_function";
constexpr StringRef kOriginalEntryFuncAttrName = "_original_entry_function";
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

    StringRef function_name = f_attr.getValue();
    // TODO(b/228928859): Improve the getter function to match attributes rather
    // than function name.
    if (!function_name.starts_with("composite_")) {
      return failure();
    }

    if (function_name.contains("depthwise_conv2d")) {
      // Uniform Quantized op requires weights of tf.DepthwiseConv2dNative to

    }

    // Only the composite functions with f32 inputs are quantizable.
    if (call_op.getResults().size() == 1 &&
        !mlir::cast<ShapedType>(call_op->getResult(0).getType())
             .getElementType()
             .isF32()) {
      check_status.Update(absl::InternalError(
          "Composite functions for quantization should be f32 type."));
    }

    // than function name.
    // If enable_legacy_weight_only is enabled, QuantizeFunctionsPattern
    // does not get called and function remains as composite
    if (!function_name.starts_with("quantized_") &&
        !function_name.starts_with("composite_")) {
      return failure();
    }

    if (function_name.contains("depthwise_conv2d")) {
      return addReshapeOpToDepthwiseWeight(call_op, rewriter);

    return false;
  }

  // Checks if op calls a composite function and all the inputs are quantized.
  bool IsCompositeFunction(TF::PartitionedCallOp call_op) const {
    if (!call_op->hasAttr(kQuantTraitAttrName)) {
      return false;
    }

    const auto f_attr = call_op.getFAttr().dyn_cast<FlatSymbolRefAttr>();
    if (!f_attr || !f_attr.getValue().starts_with("composite_")) {
      return false;
    }

"""Op composition registration."""


# TODO(fengliuai): add the tf_export decrator
class Composite(object):
  """A decorator to register a function as a composition for an TF operator.

  The argument to the decorator must be the name of a TF raw operator the
  function composites for. Decorated function must take positional arguments
  which corresponds to the input and attributes in OpDef of the TF operation.