| 208 | ], |
| 209 | ) |
| 210 | def test_dump(dump_format): |
| 211 | @trace(symbolic=True, capture_as_const=True) |
| 212 | def f(a, b): |
| 213 | return a + b |
| 214 | |
| 215 | # prevent from remaining scope from exception test |
| 216 | AutoNaming.clear() |
| 217 | a = tensor([2]) |
| 218 | b = tensor([4]) |
| 219 | y = f(a, b).numpy() |
| 220 | |
| 221 | for i in range(3): |
| 222 | np.testing.assert_equal(f(a, b).numpy(), y) |
| 223 | |
| 224 | file = io.BytesIO() |
| 225 | dump_info = f.dump(file, dump_format=dump_format) |
| 226 | assert dump_info.nr_opr == 3 |
| 227 | np.testing.assert_equal(dump_info.inputs, ["arg_0", "arg_1"]) |
| 228 | np.testing.assert_equal(dump_info.outputs, ["ADD"]) |
| 229 | file.seek(0) |
| 230 | infer_cg = cgtools.GraphInference(file) |
| 231 | result = list((infer_cg.run(a, b)).values())[0] |
| 232 | np.testing.assert_equal(result[0], y) |
| 233 | |
| 234 | |
| 235 | def test_capture_dump(): |