if (!call_op.getF().getRootReference().getValue().starts_with(
            kNestedModule))
      return WalkResult::advance();
    // This is a call we need to inline!
    LLVM_DEBUG(llvm::dbgs()
               << "Found call to inline: " << *call_op.getOperation() << "\n");

    auto call_interface = cast<CallOpInterface>(call_op.getOperation());
    auto called_func =

 private:
  LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_") ||
        !call_op->hasAttr(kQuantTraitAttrName)) {
      return failure();
    }

    return GetInverseScalesConstantOp().getValue();
  }

  func::CallOp GetCallOp() { return call_op_; }

  FlatSymbolRefAttr GetFunction() { return call_op_.getCalleeAttr(); }

 private:
  explicit UniformQuantizeFunctionCallPattern(func::CallOp call_op)
      : call_op_(call_op) {}

  func::CallOp call_op_;
};

// Matches the following pattern that represents uniform dequantization.
//

      aggregated_cost_with_decisions;
};

// If the output is produced by a callop, will return the callop, otherwise,
// will return nullptr.
inline func::CallOp GetProducerCallOpOrNull(Value output) {
  Operation* output_op = output.getDefiningOp();
  if (output_op != nullptr && llvm::isa<func::CallOp>(output_op)) {
    return llvm::cast<func::CallOp>(output_op);
  }
  return nullptr;
}

class PickSubgraphsPass

  bool changed = false;
  auto walk_result = func.walk([&](CallOp call_op) {
    auto callee = table.lookup<TFRFuncOp>(call_op.getCallee());
    if (!callee || callee.isExternal()) return WalkResult::advance();

    // Record the boundary of the inlined operations. The inlined operation will
    // be inserted between these two operations.
    Operation* inlined_point = call_op.getOperation();
    Operation* after_inlined_point =

      int current_num_elements =
          mlir::cast<ShapedType>(call_op.getOperand(idx).getType())
              .getNumElements();
      if (current_num_elements < min_num_elements_for_weights_) {
        call_op.emitRemark("Quantization is skipped for ")
            << call_op->getName().getStringRef().str() << " because it has "
            << current_num_elements

                                      TF_Status* status) {
  tensorflow::ImmediateOpPtr call_op;
  absl::Span<tensorflow::AbstractTensorHandle* const> input_span(
      reinterpret_cast<tensorflow::AbstractTensorHandle**>(
          tensorflow::unwrap(inputs)),
      static_cast<size_t>(num_inputs));
  status->status = tensorflow::unwrap(func)->MakeCallOp(input_span, &call_op);
  if (!status->status.ok()) {
    return nullptr;
  }

  tensorflow::ImmediateOpPtr call_op;
  absl::Span<tensorflow::AbstractTensorHandle* const> input_span(
      reinterpret_cast<tensorflow::AbstractTensorHandle**>(
          tensorflow::unwrap(inputs)),
      static_cast<size_t>(num_inputs));
  status->status = tensorflow::unwrap(func)->MakeCallOp(input_span, &call_op);
  if (!status->status.ok()) {
    return nullptr;
  }

    if (method.ok() && method->has_static_range_ptq()) return true;
  }

  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(

  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
  EXPECT_TRUE(HasWeightOnlyPtqMethod(call_op));
}

TEST_F(LiftAsFunctionCallTest, HasWeightOnlyPtqMethodDifferentMethod) {
  const absl::string_view kModuleDotNoQuantization = R"mlir(
    module {