Repository: pyg-team/pytorch_geometric
Pull request #7221 (2023-04-22 11:22:12)

Test `return_emb` in `MLP`
There was a bug in the .forward method of the MLP class. If return_emb=False, then it would still return a tuple containing the embeddings. The line with the bug was:

```python
return (x, emb) if isinstance(return_emb, bool) else x
```
since `isinstance(False, bool)` is True. 


codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7221?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7221](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7221?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4f7d650) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/a67799fb8688175ebedbdcf8365de5b62c83859c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (a67799f) will **decrease** coverage by `0.44%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7221      +/-   ##
==========================================
- Coverage   92.03%   91.59%   -0.44%     
==========================================
  Files         437      437              
  Lines       24092    24089       -3     
==========================================
- Hits        22172    22064     -108     
- Misses       1920     2025     +105     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7221?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/models/mlp.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7221?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9tbHAucHk=) | `95.18% <100.00%> (-1.21%)` | :arrow_down: |

... and [20 files with indirect coverage changes](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7221/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/models/mlp.py b/torch_geometric/nn/models/mlp.py
index 402a6c087..d2be2475e 100644
--- a/torch_geometric/nn/models/mlp.py
+++ b/torch_geometric/nn/models/mlp.py
@@ -194,9 +194,14 @@ class MLP(torch.nn.Module):
         Args:
             x (torch.Tensor): The source tensor.
             return_emb (bool, optional): If set to :obj:`True`, will
-                additionally return the embeddings before execution of to the
+                additionally return the embeddings before execution of the
                 final output layer. (default: :obj:`False`)
         """
+        # `return_emb` is annotated here as `NoneType` to be compatible with
+        # TorchScript, which does not support different return types based on
+        # the value of an input argument.
+        emb: Optional[Tensor] = None
+
         for i, (lin, norm) in enumerate(zip(self.lins, self.norms)):
             x = lin(x)
             if self.act is not None and self.act_first:
@@ -205,7 +210,8 @@ class MLP(torch.nn.Module):
             if self.act is not None and not self.act_first:
                 x = self.act(x)
             x = F.dropout(x, p=self.dropout[i], training=self.training)
-            emb = x
+            if isinstance(return_emb, bool) and return_emb is True:
+                emb = x
 
         if self.plain_last:
             x = self.lins[-1](x)


Test patch:
diff --git a/test/nn/models/test_mlp.py b/test/nn/models/test_mlp.py
index cb28a89ac..d830a8ed7 100644
--- a/test/nn/models/test_mlp.py
+++ b/test/nn/models/test_mlp.py
@@ -2,7 +2,6 @@ import pytest
 import torch
 
 from torch_geometric.nn import MLP
-from torch_geometric.testing import is_full_test
 
 
 @pytest.mark.parametrize('norm', ['batch_norm', None])
@@ -22,9 +21,8 @@ def test_mlp(norm, act_first, plain_last):
     out = mlp(x)
     assert out.size() == (4, 64)
 
-    if is_full_test():
-        jit = torch.jit.script(mlp)
-        assert torch.allclose(jit(x), out)
+    jit = torch.jit.script(mlp)
+    assert torch.allclose(jit(x), out)
 
     torch.manual_seed(12345)
     mlp = MLP(
@@ -39,6 +37,20 @@ def test_mlp(norm, act_first, plain_last):
     assert torch.allclose(mlp(x), out)
 
 
+def test_mlp_return_emb():
+    x = torch.randn(4, 16)
+
+    mlp = MLP([16, 32, 1])
+
+    out, emb = mlp(x, return_emb=True)
+    assert out.size() == (4, 1)
+    assert emb.size() == (4, 32)
+
+    out, emb = mlp(x, return_emb=False)
+    assert out.size() == (4, 1)
+    assert emb is None
+
+
 @pytest.mark.parametrize('plain_last', [False, True])
 def test_fine_grained_mlp(plain_last):
     mlp = MLP(