()
| 60 | |
| 61 | |
| 62 | def test_replace_var(): |
| 63 | |
| 64 | a = Tensor([1, 2]) |
| 65 | b = Tensor([3, 4]) |
| 66 | |
| 67 | @trace(symbolic=True, capture_as_const=True) |
| 68 | def fwd(a, b): |
| 69 | return (a + b) * 2 |
| 70 | |
| 71 | fwd(a, b) |
| 72 | orig_model = io.BytesIO() |
| 73 | fwd.dump( |
| 74 | orig_model, arg_names=["a", "b"], output_names="o", optimize_for_inference=False |
| 75 | ) |
| 76 | orig_model.seek(0) |
| 77 | |
| 78 | graph = Net.load(orig_model) |
| 79 | vara = graph.var_filter.name("a").as_unique() |
| 80 | varb = graph.var_filter.name("b").as_unique() |
| 81 | |
| 82 | out = F.mul(vara, varb) |
| 83 | out = F.relu(out) |
| 84 | |
| 85 | opnode = list(graph.opr_filter.has_input(vara)) |
| 86 | repl_dict = {opnode[0].outputs[0]: out} |
| 87 | graph.replace_vars(repl_dict) |
| 88 | |
| 89 | modified_model = io.BytesIO() |
| 90 | graph.dump(modified_model) |
| 91 | modified_model.seek(0) |
| 92 | load_graph = GraphInference(modified_model) |
| 93 | |
| 94 | out = load_graph.run(a, b) |
| 95 | np.testing.assert_equal(out["o"], [6, 16]) |
| 96 | |
| 97 | |
| 98 | def test_replace_opr(): |
nothing calls this directly
no test coverage detected