| 103 | |
| 104 | |
| 105 | def test_get_opr_seq(): |
| 106 | class Net(M.Module): |
| 107 | def __init__(self): |
| 108 | super().__init__() |
| 109 | self.data = megengine.tensor( |
| 110 | np.random.random((1, 1, 4, 4)), dtype=np.float32 |
| 111 | ) |
| 112 | |
| 113 | def forward(self, input): |
| 114 | A = input.shape[0] |
| 115 | shape = astensor1d((A, A), self.data, dtype="int32", device=input.device) |
| 116 | x = F.reshape(self.data, shape) |
| 117 | o = input + x |
| 118 | return o |
| 119 | |
| 120 | net = Net() |
| 121 | input = megengine.tensor(np.random.random((4, 4)), dtype=np.float32) |
| 122 | |
| 123 | @trace(symbolic=True, capture_as_const=True) |
| 124 | def func(inp, *, net=None): |
| 125 | return net(inp) |
| 126 | |
| 127 | func(input, net=net) |
| 128 | file = io.BytesIO() |
| 129 | func.dump(file, optimize_for_inference=False) |
| 130 | file.seek(0) |
| 131 | outputs = mgb_graph.load_graph(file).output_vars_list |
| 132 | |
| 133 | seq_1 = cgtools.get_oprs_seq(outputs, True) |
| 134 | assert len(seq_1) == 5 |
| 135 | |
| 136 | seq_2 = cgtools.get_oprs_seq(outputs, False) |
| 137 | assert len(seq_2) == 6 |
| 138 | |
| 139 | |
| 140 | def test_topological_sort(): |