Print relax AST in structured format. The shape of Node is ignored.
| 81 | |
| 82 | @relax.expr_functor.visitor |
| 83 | class ASTPrinter(PyExprVisitor): |
| 84 | """Print relax AST in structured format. The shape of Node is ignored.""" |
| 85 | |
| 86 | def __init__(self) -> None: |
| 87 | super().__init__() |
| 88 | self.log = ASTLog() |
| 89 | |
| 90 | def visit_constant_(self, op: Constant) -> None: |
| 91 | self.log.add("Constant") |
| 92 | |
| 93 | def visit_global_var_(self, op: GlobalVar) -> None: |
| 94 | self.log.add("GlobalVar") |
| 95 | |
| 96 | def visit_tuple_(self, op: Tuple) -> None: |
| 97 | self.log.add("Tuple") |
| 98 | self.log.push_scope() |
| 99 | for field in op.fields: |
| 100 | self.visit_expr(field) |
| 101 | self.log.pop_scope() |
| 102 | |
| 103 | def visit_var_(self, op: Var) -> None: |
| 104 | self.log.add("Var") |
| 105 | |
| 106 | def visit_dataflow_var_(self, op: DataflowVar) -> None: |
| 107 | self.log.add("DataflowVar") |
| 108 | |
| 109 | def visit_function_(self, op: Function) -> None: |
| 110 | self.log.add("Function") |
| 111 | self.log.push_scope() |
| 112 | for param in op.params: |
| 113 | self.visit_var_def(param) |
| 114 | |
| 115 | self.visit_expr(op.body) |
| 116 | self.log.pop_scope() |
| 117 | |
| 118 | def visit_call_(self, op: Call) -> None: |
| 119 | self.log.add("Call") |
| 120 | self.log.push_scope() |
| 121 | self.visit_expr(op.op) |
| 122 | |
| 123 | for arg in op.args: |
| 124 | self.visit_expr(arg) |
| 125 | self.log.pop_scope() |
| 126 | |
| 127 | def visit_if_(self, op: If) -> None: |
| 128 | self.log.add("If") |
| 129 | self.log.push_scope() |
| 130 | self.visit_expr(op.cond) |
| 131 | self.visit_expr(op.true_branch) |
| 132 | self.visit_expr(op.false_branch) |
| 133 | self.log.pop_scope() |
| 134 | |
| 135 | def visit_op_(self, op: Op) -> None: |
| 136 | self.log.add("Op") |
| 137 | |
| 138 | def visit_tuple_getitem_(self, op: TupleGetItem) -> None: |
| 139 | self.log.add("TupleGetItem") |
| 140 | self.log.push_scope() |
no outgoing calls
searching dependent graphs…