{DT_FLOAT}, f_name_attr);

  TF_ASSERT_OK(root.ToGraph(&graph));
  EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}

TEST(IsInferenceGraph, GraphDoesNotContrainTPUPartitionedCall) {
  FunctionDef fd = FunctionDefHelper::Define(
      // Name
      "XTimesTwoFloat",
      // Args
      {"x: float"},

// Check if a graph contains TPUPartitionedCall op, including its reachable
// functions. The function library is used to store the functions that are
// defined in a TensorFlow program
bool IsInferenceGraph(const Graph& graph,
                      const FunctionLibraryDefinition* function_library);
}  // namespace tensorflow

    func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {
      has_tpu_partitioned_call = true;
    });
    if (has_tpu_partitioned_call) break;
  }
  return has_tpu_partitioned_call;
}

bool IsInferenceGraph(const Graph& graph,
                      const FunctionLibraryDefinition* function_library) {
  if (DoesGraphContainTPUPartitionedCall(graph)) return true;
  GraphDef graph_def;
  graph.ToGraphDef(&graph_def);