xref: /aosp_15_r20/external/tensorflow/tensorflow/python/grappler/graph_analyzer.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# =============================================================================
15"""A tool that finds all subgraphs of a given size in a TF graph.
16
17The subgraph patterns are sorted by occurrence, and only the transitive fanin
18part of the graph with regard to the fetch nodes is considered.
19"""
20
21import argparse
22import sys
23
24from absl import app
25
26from tensorflow.python.grappler import _pywrap_graph_analyzer as tf_wrap
27
28
29def main(_):
30  tf_wrap.GraphAnalyzer(FLAGS.input, FLAGS.n)
31
32
33if __name__ == "__main__":
34  parser = argparse.ArgumentParser()
35  parser.add_argument(
36      "--input",
37      type=str,
38      default=None,
39      help="Input file path for a TensorFlow MetaGraphDef.")
40  parser.add_argument(
41      "--n", type=int, default=None, help="The size of the subgraphs.")
42  FLAGS, unparsed = parser.parse_known_args()
43  app.run(main=main, argv=[sys.argv[0]] + unparsed)
44