r"""Returns :class:`.tensor.core.megbrain_graph.VarNode` of type ``var_type`` that input ``var`` depands on. If ``var_type`` is None, returns all types.
(
var: Union[_VarNode, List[_VarNode]], var_type: Union[str, List[str]] = None
)
| 28 | |
| 29 | |
| 30 | def get_dep_vars( |
| 31 | var: Union[_VarNode, List[_VarNode]], var_type: Union[str, List[str]] = None |
| 32 | ) -> List[_VarNode]: |
| 33 | r"""Returns :class:`.tensor.core.megbrain_graph.VarNode` of type ``var_type`` that input ``var`` |
| 34 | depands on. If ``var_type`` is None, returns all types. |
| 35 | """ |
| 36 | outputs = [] |
| 37 | memo = set() |
| 38 | |
| 39 | if isinstance(var, _VarNode): |
| 40 | var = [var] |
| 41 | |
| 42 | if isinstance(var_type, str): |
| 43 | var_type = [var_type] |
| 44 | |
| 45 | q = list(var) |
| 46 | while q: |
| 47 | v = q.pop(0) |
| 48 | if v in memo: |
| 49 | continue |
| 50 | memo.add(v) |
| 51 | q.extend(get_owner_opr_inputs(v)) |
| 52 | if var_type is not None: |
| 53 | if get_owner_opr_type(v) in var_type: |
| 54 | outputs.append(v) |
| 55 | else: |
| 56 | outputs.append(v) |
| 57 | |
| 58 | return outputs |
| 59 | |
| 60 | |
| 61 | def get_owner_opr_inputs(var: _VarNode) -> List[_VarNode]: |
no test coverage detected