// function.
  SmallVector<TF::XlaCallModuleOp> xla_call_module_ops;
  for (auto xla_call_module_op : func_ops[0].getOps<TF::XlaCallModuleOp>()) {
    xla_call_module_ops.push_back(xla_call_module_op);
  }
  ASSERT_THAT(xla_call_module_ops, SizeIs(1));
  auto xla_call_module_op = xla_call_module_ops[0];
  EXPECT_THAT(xla_call_module_op,
              HasStringAttr("_tfl_quant_trait", StrEq("fully_quantizable")));

  ASSERT_THAT(main_fn, NotNull());

  auto xla_call_module_ops = main_fn.getOps<TF::XlaCallModuleOp>();
  ASSERT_FALSE(xla_call_module_ops.empty());

  // Test that `GetQuantizationMethod` returns a valid `Method` corresponding to
  // `"no_quantization {}"`.
  const absl::StatusOr<Method> method =
      GetQuantizationMethod(*xla_call_module_ops.begin());
  ASSERT_THAT(method, IsOk());

  ASSERT_TRUE(module_op);

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

  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  ASSERT_THAT(xla_call_module_op, NotNull());

  FlatSymbolRefAttr xla_call_op_attr = GetFuncAttr(xla_call_module_op);
  EXPECT_EQ(xla_call_op_attr.getValue(), "composite_fn_1");
}

TEST_F(AttrsAndConstraintsTest, PartitionedCallGetFuncAttr) {

  ASSERT_TRUE(module_op);

  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
  ASSERT_TRUE(succeeded(xla_call_module_op));

  const std::unique_ptr<OpQuantSpec> op_quant_spec =
      GetStableHloOpQuantSpec(*xla_call_module_op);
  ASSERT_THAT(op_quant_spec, NotNull());

  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim, IsEmpty());

  return result;
}

// Retrieves `QuantizationResult` from `xla_call_module_op`. If
// `xla_call_module_op` is a quantizable unit, then a `QuantizationResult` will
// be returned with its `name` field set to the callee's name. The `method`
// field will be set to `NoQuantization` because remaining `xla_call_module_op`s
// means they are not quantized. Returns `std::nullopt` if `xla_call_module_op`
// is not a quantizable unit.

bool IsQuantizedXlaCallModuleOp(TF::XlaCallModuleOp xla_call_module_op) {
  return !IsOpNotQuantized(xla_call_module_op) &&
         xla_call_module_op->hasAttr(kQuantTraitAttrName) &&
         xla_call_module_op->hasAttr(kEntryFuncAttrName);
}

// Returns the entry function, i.e. the callee of `xla_call_module_op`.
func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op,

      : match_regex_(GetMatchRegex(spec)) {}

  // Returns `true` when matched with the entry function of
  // `xla_call_module_op`.
  bool Match(TF::XlaCallModuleOp xla_call_module_op) const {
    if (match_regex_ == nullptr) return false;

    const std::string lifted_func_name =
        xla_call_module_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function")
            .getValue()
            .str();

  if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
      xla_call_module_op != nullptr) {
    absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
    if (method.ok() && method->has_no_quantization()) {
      return true;
    }
  }
  return false;
}

// Populates `spec.coeff_op_quant_dim` according to `xla_call_module_op`'s

bool IsCallToQuantizableLiftedFunction(Operation *op) {
  if (!op) return false;
  if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
      xla_call_module_op != nullptr) {
    absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
    if (method.ok() && method->has_static_range_ptq()) return true;
  }

      /*infer_tensor_ranges=*/true, /*legacy_float_scale=*/false,
      /*is_qdq_conversion=*/false);
  Operation* xla_call_module_op =
      FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  Operation* filter_dcast_op =
      xla_call_module_op->getOperand(1).getDefiningOp();
  Operation* filter_qcast_op = filter_dcast_op->getOperand(0).getDefiningOp();
  ASSERT_NE(filter_qcast_op, nullptr);