| 127 | |
| 128 | @pytest.mark.parametrize("trace_mode", [False, True]) |
| 129 | def test_elemwise_fuse(trace_mode): |
| 130 | # explicitly declare opt_level as 2 |
| 131 | @trace(symbolic=trace_mode, opt_level=2) |
| 132 | def f(a, b): |
| 133 | base = 0 |
| 134 | c = b - a |
| 135 | _, idx = F.topk(c, 3) |
| 136 | # internally, biased_idx will be idx as gopt will ignore the addition |
| 137 | biased_idx = base + idx |
| 138 | return biased_idx |
| 139 | |
| 140 | a = tensor(np.ones((7, 2)), dtype=np.int32) |
| 141 | b = tensor(2 * np.ones((7, 2)), dtype=np.float32) |
| 142 | |
| 143 | for i in range(3): |
| 144 | y = f(a, b) |
| 145 | y.numpy() |
| 146 | |
| 147 | |
| 148 | @pytest.mark.parametrize("trace_mode", [False, True]) |