(fuse_dimshuffle, fuse_reduce)
| 732 | @pytest.mark.require_ngpu(1) # nvrtc backend |
| 733 | def test_trace_jit_config(): |
| 734 | def run(fuse_dimshuffle, fuse_reduce): |
| 735 | config = GraphOptimizationConfig() |
| 736 | config.jit_fuse_dimshuffle = fuse_dimshuffle |
| 737 | config.jit_fuse_reduce = fuse_reduce |
| 738 | |
| 739 | # set opt_level = 1 to avoid fusing dimshuffle and reduce at the same time |
| 740 | @trace(opt_level=1, graph_opt_config=config) |
| 741 | def func(x): |
| 742 | return x + 1 |
| 743 | |
| 744 | x = tensor(2) |
| 745 | y = func(x) |
| 746 | y = func(x) |
| 747 | # func._compile() |
| 748 | |
| 749 | options = func._trace.options |
| 750 | mapping = {None: 0, False: 1, True: 2} |
| 751 | assert options.graph_opt.jit == 0 |
| 752 | assert options.graph_opt.jit_config.fuse_dimshuffle == mapping[fuse_dimshuffle] |
| 753 | assert options.graph_opt.jit_config.fuse_reduce == mapping[fuse_reduce] |
| 754 | |
| 755 | for fuse_dimshuffle in [None, False, True]: |
| 756 | for fuse_reduce in [None, False, True]: |
no test coverage detected