| 42 | |
| 43 | |
| 44 | class OpWrapper: |
| 45 | def __init__(self, op, graph): |
| 46 | assert isinstance(graph, GraphWrapper) |
| 47 | self._op = op |
| 48 | self._graph = graph |
| 49 | |
| 50 | def type(self): |
| 51 | """ |
| 52 | Get the type of this operator. |
| 53 | """ |
| 54 | return self._op.type |
| 55 | |
| 56 | def inputs(self, name): |
| 57 | """ |
| 58 | Get all the variables by the input name. |
| 59 | """ |
| 60 | if name in self._op.input_names: |
| 61 | return [ |
| 62 | self._graph.var(var_name) for var_name in self._op.input(name) |
| 63 | ] |
| 64 | else: |
| 65 | return [] |
| 66 | |
| 67 | def outputs(self, name): |
| 68 | """ |
| 69 | Get all the variables by the output name. |
| 70 | """ |
| 71 | return [self._graph.var(var_name) for var_name in self._op.output(name)] |
| 72 | |
| 73 | |
| 74 | class GraphWrapper: |