// Checks if an op calls a composite function and all the inputs and outputs are
// quantized.
bool IsQuantizedCompositeFunction(func::CallOp call_op) {
  if (!call_op.getCallee().starts_with("quantized_")) {
    return false;
  }

  bool has_quantized_types = false;
  for (Value operand : call_op.getOperands()) {
    if (const TensorType type = mlir::dyn_cast<TensorType>(operand.getType())) {

// quantization.
inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
}

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}

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

  Operation* call_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  EXPECT_FALSE(IsHybridQuantizedOp(call_op));
}

constexpr absl::string_view kModuleDotGeneralFullyConnected = R"mlir(
  module {

    if (!main_func) continue;

    SymbolTable symbol_table(module_op);
    for (auto call_op : main_func.getOps<TF::PartitionedCallOp>()) {
      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
    }
    for (auto call_op : main_func.getOps<TF::StatefulPartitionedCallOp>()) {
      func_ops.push_back(

          Operation* next_op = cur_op_use.getOwner();
          int next_op_operand_num = cur_op_use.getOperandNumber();
          if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(next_op)) {
            mlir::func::FuncOp func =
                llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
            if (!func) continue;
            next_values_to_visit.push_back(

           /*dep=*/{"assert"}},
      });
  TF_ASSERT_OK(ContextFromInterface(unwrap(context.get()))
                   ->AddFunctionDef(assert_and_collective));

  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> call_op(
      TFE_NewOp(context.get(), "AssertAndCollective", status.get()),
      TFE_DeleteOp);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
  std::unique_ptr<ParallelTensor> reduced_values =

    result_idx += func_op.getNumResults();

    auto call_op = builder.create<TF::PartitionedCallOp>(
        module_op.getLoc(), new_types, new_args,
        SymbolRefAttr::get(context, func_op.getSymName()),
        /*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));
    call_op_returns.append(call_op.getResults().begin(),

      return failure();
    }
  }

  return success();
}

LogicalResult CheckOutputConsumer(
    Operation* call_op, int expected_num_outputs,
    llvm::DenseSet<int> expected_consumer_indices) {
  const int num_results = call_op->getNumResults();
  if (num_results != expected_num_outputs) return failure();

  for (int i = 0; i < expected_num_outputs; ++i) {

  }
  auto a = Parameter(&builder, 0, ShapeUtil::MakeScalarShape(xla::F32), "a");
  auto b = Parameter(&builder, 1, ShapeUtil::MakeScalarShape(xla::F32), "b");
  XlaOp call_op = xla::Call(&builder, to_apply, {a, b});

  std::vector<XlaOp> tuple_values;
  tuple_values.push_back(call_op);
  xla::Tuple(&builder, tuple_values);

  TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build());

          next_values_to_visit.push_back(use.getOwner()->getResult(0));
          continue;
        }

        if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(owner)) {
          mlir::func::FuncOp func =
              llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
          if (!func) continue;
          next_values_to_visit.push_back(
              func.getArgument(use.getOperandNumber()));