ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  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 =

      ParseModuleOpString(kModuleDotGeneralFullyConnected);
  ASSERT_TRUE(module_op);

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

  auto dot_general_op = *main_fn.getOps<DotGeneralOp>().begin();
  EXPECT_THAT(IsDotGeneralFullyConnected(dot_general_op), true);
}

TEST_F(AttrsAndConstraintsTest, IsDotGeneralFullyConnectedReturnsFalse) {
  OwningOpRef<ModuleOp> module_op =

  llvm::SmallVector<TF::VarHandleOp, 4> variables;
  // Capture list of all read only variables.
  for (auto func : module.getOps<func::FuncOp>()) {
    if (func == session_init_func) continue;
    for (auto var_handle_op : func.getOps<TF::VarHandleOp>()) {
      if (!analyzer.IsPotentiallyWritten(var_handle_op.getResource())) {
        variables.push_back(var_handle_op);
      }
    }
  }

    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(

  bool need_prefix_for_output_name = false;
  std::vector<StringRef> fn_input_name_vec, fn_output_name_vec;
  StringSet<> input_name_set, output_name_set;
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (!ShouldIncludeInMainFunction(func_op)) continue;
    if (auto tf_attrs =
            func_op->getAttrOfType<DictionaryAttr>(kEntryFunctionAttr)) {

template <typename GemmStyleOp>
// Match for all gemm_style op and check for possible fusions.
LogicalResult MatchGemmStyleOp(func::FuncOp entry_func_op) {
  const auto op_iterator_range = entry_func_op.getOps<GemmStyleOp>();
  if (op_iterator_range.empty()) {
    LLVM_DEBUG(llvm::dbgs() << "Function does not have "
                            << GemmStyleOp::getOperationName() << " op.\n");
    return failure();
  }

      }
    }
    Builder stablehlo_builder(stablehlo_module_op.get().getContext());
    // Rename XlaCallModuleOp's functions to avoid naming conflicts.
    for (auto func_op : stablehlo_module_op.get().getOps<func::FuncOp>()) {
      const std::string new_func_name =
          CreateNewFuncName(func_op.getSymName(), parent_module_symbol_table);
      if (failed(stablehlo_module_symbol_table.replaceAllSymbolUses(

  // to store its `SymbolTable` to `SymbolTableCollection`.
  SymbolTable stablehlo_symbol_table(stablehlo_module);

  Builder builder(context);
  StringAttr main_func_name;
  for (auto func : stablehlo_module.getOps<func::FuncOp>()) {
    const bool is_main_func = func.getSymName() == kStablehloMainFunctionName;
    if (tf_symbol_table.lookup(func.getSymName())) {
      if (failed(stablehlo_symbol_table.renameToUnique(

    if (function_name.contains("conv")) {
      return (*(func.getOps<mlir::stablehlo::ConvolutionOp>().begin()))
          .getDimensionNumbers()
          .getKernelOutputFeatureDimension();
    } else if (function_name.contains("dot_general")) {
      auto dot = *(func.getOps<mlir::stablehlo::DotGeneralOp>().begin());
      const ::mlir::stablehlo::DotDimensionNumbersAttr dimension_numbers =

}  // namespace

void LegalizeTFGToTFE::runOnOperation() {
  MLIRContext &context = getContext();
  ModuleOp module = getOperation();

  DenseSet<StringRef> func_symbols;
  for (auto &op : module.getBodyRegion().getOps()) {
    if (auto func = llvm::dyn_cast<tfg::GraphFuncOp>(op)) {
      func_symbols.insert(
          func->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
              .getValue());
    }
  }