File tree Expand file tree Collapse file tree 1 file changed +23
-0
lines changed Expand file tree Collapse file tree 1 file changed +23
-0
lines changed Original file line number Diff line number Diff line change
1
+ # Copyright 2025 Arm Limited and/or its affiliates.
2
+ #
3
+ # This source code is licensed under the BSD-style license found in the
4
+ # LICENSE file in the root directory of this source tree.
5
+
6
+ import torch
7
+ from executorch .devtools .visualization .visualization_utils import visualize_graph
8
+ from executorch .exir import ExportedProgram
9
+ from executorch .exir .pass_base import ExportPass , PassResult
10
+
11
+
12
+ class VisualizePass (ExportPass ):
13
+ """
14
+ This pass visualizes the graph at the point of insertion in the pass manager
15
+ """
16
+
17
+ def __init__ (self , exported_program : ExportedProgram ) -> None :
18
+ super ().__init__ ()
19
+ self .exported_program = exported_program
20
+
21
+ def call (self , graph_module : torch .fx .GraphModule ) -> PassResult :
22
+ visualize_graph (graph_module , self .exported_program )
23
+ return PassResult (graph_module , False )
You can’t perform that action at this time.
0 commit comments