function_name.contains("conv3d"))) {
      call_op->removeAttr(kQuantTraitAttrName);
    }

    // TODO(b/270906404): Support weight-only gather for uniform quantized opset
    // in PTQ mode
    if (target_opset_ == OpSet::UNIFORM_QUANTIZED &&
        function_name.contains("gather")) {
      call_op->removeAttr(kQuantTraitAttrName);
    }

    return failure();
  }
  QuantMethod quantization_method_;

                       const StringAttr &new_ref_func_name) {
  // Create PartitionedCallOp to the copied composite function. This
  // PartitionedCallOp does not have kQuantTraitAttrName, and therefore won't
  // get quantized.
  auto new_call_op = rewriter.create<TF::PartitionedCallOp>(
      call_op.getLoc(), call_op.getResultTypes(), call_op.getOperands(),

  }

  TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
  return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
             llvm::StringRef(
                 QuantTraitValues[QuantizationTrait::FullyQuantizable]);
}

// Returns the composite function name.

  if (GetStableHloQuantConstraints(op)->has_same_scale_requirement) {
    return true;
  }

  const bool attr_enforced_quantizable =
      op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
      op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue().str() ==
          QuantTraitValues[QuantizationTrait::FullyQuantizable];
  return attr_enforced_quantizable;
}

// Checks whether the given op contains QuantizationTrait::FullyQuantizable.
inline bool HasQuantizableTrait(Operation* op) {
  return op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue().str() ==
             QuantTraitValues[QuantizationTrait::FullyQuantizable];
}

// Returns true if `op` has two operands and one result and only second operand

    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_") ||
        !call_op->hasAttr(kQuantTraitAttrName)) {
      return failure();
    }

    absl::Status check_status;
    // TODO(b/270906404): Support weight-only gather for uniform quantized opset
    // in PTQ mode

                           SymbolTable& symbol_table) {
  // Do not inline lifted quantized functions used for fusing patterns.
  // TODO - b/310539922: Remove reference to TF/TFL utils.
  if (call_op->hasAttr(kQuantTraitAttrName)) {
    return;
  }

  auto function_name = call_op
                           ->getAttrOfType<FlatSymbolRefAttr>(
                               TF::kStablehloEntryFunctionAttrName)

      /*config=*/"", /*config_proto=*/"", /*executor_type=*/"");

  // Set the attribute to annotate this function call op as a quantizable spot.
  call_op->setAttr(
      kQuantTraitAttrName,
      builder.getStringAttr(StringRef(
          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));

  return call_op.getOutput();
}

                                PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
    // Non-quantizable op
    if (!op->hasAttr(kQuantTraitAttrName)) return failure();
    StringRef function_name = f_attr.getValue();
    // TODO(b/228928859): Improve the getter function to match attributes rather
    // than function name.

// 2. `xla_call_module_op` should have the `kQuantTraitAttrName` attribute.
// 3. It should also have the `kEntryFuncAttrName` attribute, which points to
//    the function that `xla_call_module_op` represents.
bool IsQuantizedXlaCallModuleOp(TF::XlaCallModuleOp xla_call_module_op) {
  return !IsOpNotQuantized(xla_call_module_op) &&
         xla_call_module_op->hasAttr(kQuantTraitAttrName) &&