Repository: pyg-team/pytorch_geometric
Pull request #7375 (2023-05-16 18:25:04)

Use recommended `script_if_tracing` in library code

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7375](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (3e46cbb) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/239c4b5a8d7b9862ba942c6482cb2d5077533475?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (239c4b5) will **decrease** coverage by `0.40%`.
> The diff coverage is `50.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7375      +/-   ##
==========================================
- Coverage   91.75%   91.35%   -0.40%     
==========================================
  Files         438      438              
  Lines       24438    24435       -3     
==========================================
- Hits        22422    22323      -99     
- Misses       2016     2112      +96     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/pool/connect/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvY29ubmVjdC9iYXNlLnB5) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/nn/pool/select/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2VsZWN0L2Jhc2UucHk=) | `96.42% <100.00%> (ø)` | |

... and [20 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7375/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/torch_geometric/nn/pool/connect/base.py b/torch_geometric/nn/pool/connect/base.py
index fd78a5c61..262937603 100644
--- a/torch_geometric/nn/pool/connect/base.py
+++ b/torch_geometric/nn/pool/connect/base.py
@@ -7,7 +7,7 @@ from torch import Tensor
 from torch_geometric.nn.pool.select import SelectOutput
 
 
-@torch.jit.script
+@torch.jit.script_if_tracing
 @dataclass(init=False)
 class ConnectOutput:
     r"""The output of the :class:`Connect` method, which holds the coarsened
diff --git a/torch_geometric/nn/pool/select/base.py b/torch_geometric/nn/pool/select/base.py
index fcde48234..ce9281653 100644
--- a/torch_geometric/nn/pool/select/base.py
+++ b/torch_geometric/nn/pool/select/base.py
@@ -5,7 +5,7 @@ import torch
 from torch import Tensor
 
 
-@torch.jit.script
+@torch.jit.script_if_tracing
 @dataclass(init=False)
 class SelectOutput:
     r"""The output of the :class:`Select` method, which holds an assignment


Test patch:
diff --git a/test/nn/pool/select/test_select_topk.py b/test/nn/pool/select/test_select_topk.py
index a0418fff9..571eaa57a 100644
--- a/test/nn/pool/select/test_select_topk.py
+++ b/test/nn/pool/select/test_select_topk.py
@@ -47,7 +47,7 @@ def test_select_topk(min_score):
         assert str(pool) == 'SelectTopK(16, min_score=2.0)'
 
     out = pool(x, batch)
-    assert isinstance(out, SelectOutput)
+    assert isinstance(out, SelectOutput.__original_fn)
 
     assert out.num_nodes == 6
     assert out.num_clusters <= out.num_nodes