Fix `fx` problem for `Aggregation` #5021

Padarn · 2022-07-21T02:06:38Z

This PR addresses a problem we had when using torch.fx on our new nn.aggr.Aggregation class. This class redefines the __call__ function to contain:

     def __call__(self, x: Tensor, index: Optional[Tensor] = None,
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2, **kwargs) -> Tensor:

        if dim >= x.dim() or dim < -x.dim():
            raise ValueError(f"Encountered invalid dimension '{dim}' of "
                             f"source tensor with {x.dim()} dimensions")
        ...

During a torch.fx._symbolic_trace pytorch modules torch.nn.Module calls are excluded from symbolic trancing intentionally using is_leaf_module; this allows for avoiding errors on symbolic tracing of for example, conditions statements involving tensors (which are not handled).

However, this is implemented by patching a specific function within torch.fx._symbolic_trace:

patcher.patch_method(torch.nn.Module, "__call__", module_call_wrapper, deduplicate=False)

from here.

This doesn't patch our __call__ as it no longer lives in torch.nn.Module.

The fix here is to reimplemnt the trace function to add an extra wrapper around our code.

Note that this should be replaced by a better solution once one is available, by for example:

Allowing for kwargs in pre_forward_hooks: nn.Module hooks ignore kwargs pytorch/pytorch#35643 (comment)
Getting a better way to extend the patched functions from the torch.fx team.

rusty1s · 2022-07-21T05:30:08Z

torch_geometric/nn/fx.py

    if op == 'call_module':
-        return isinstance(get_submodule(module, target), GlobalPooling)
+        return isinstance(get_submodule(module, target),
+                          GlobalPooling) or isinstance(


Let‘s just replace it here.

rusty1s · 2022-07-21T05:31:13Z

torch_geometric/nn/aggr/base.py


    # @abstractmethod
    def forward(self, x: Tensor, index: Optional[Tensor] = None,
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
-                dim: int = -2) -> Tensor:
+                dim: int = -2, **kwargs) -> Tensor:


Why do we need to add kwargs here in the first place?

because of the GraphMultisetTransformer - which also takes edge_index optionally

If we can just add this to the default args then the problem is solved, but I don't see it as a common need

rusty1s · 2022-07-21T05:33:10Z

torch_geometric/nn/aggr/base.py

+    def __init__(self):
+        super().__init__()
+        self._forward_sub = self.forward
+        self.forward = self._foward


I think this is overly complicated. How about we just do

def forward: self.validate() self._forward()

And overwrite _forward in child modules?

yeah I was thinking that too - just providing this option incase we want to keep the current abstract interface

rusty1s · 2022-07-21T05:35:05Z

torch_geometric/nn/aggr/base.py

@@ -35,35 +75,6 @@ def forward(self, x: Tensor, index: Optional[Tensor] = None,
    def reset_parameters(self):
        pass

-    def __call__(self, x: Tensor, index: Optional[Tensor] = None,


Just so I understand better: what was the issue with this code to begin with? It looks like the tracer has problems when overwriting call. Is that correct? Is there any way to fix this in the tracer?

yes, its because of this: https://github.com/pytorch/pytorch/blob/e68583b4d180066b8e4f108e0d23176a2676421c/torch/fx/_symbolic_trace.py#L702
the tracer only looks for a leaf module in the patched methods - but since we overwrite call its no longer patched (the super().__call__ is still patched)

Can we patch it ourselves?

hmm, yes... but we'll have to copy quite a bit of code from their implementation. I can implement that and we can see what it looks like?

Ok, depends on how easy it would be to try out and how large the code to patch would be. Do not burn too much time into it if it is not easy to fix :)

I thought about this, but I don't know how the _Patcher context manager works, so I was afraid to suggest it even if it happened to pass tests.

Could certainly invest time in understanding if you think its a big upgrade thouh.

Can we quickly test if it works? I am always in favor of reducing number of lines :)

Hmm, no doesn't work 'out of the box'.

Ok, then leave it as it is. Might be good to add some exhaustive documentation here, and also link to the PyTorch file this code is coming from.

Agreed. Will keep an eye open for better solutions too.

codecov · 2022-07-22T02:50:20Z

Codecov Report

Merging #5021 (6827cd9) into master (be9e4af) will decrease coverage by 1.94%.
The diff coverage is 93.61%.

@@            Coverage Diff             @@
##           master    #5021      +/-   ##
==========================================
- Coverage   84.77%   82.83%   -1.95%     
==========================================
  Files         331      331              
  Lines       18115    18156      +41     
==========================================
- Hits        15357    15039     -318     
- Misses       2758     3117     +359

Impacted Files	Coverage Δ
torch_geometric/nn/to_hetero_transformer.py	`95.26% <ø> (ø)`
torch_geometric/nn/fx.py	`90.05% <93.61%> (+1.01%)`	⬆️
torch_geometric/nn/models/dimenet_utils.py	`0.00% <0.00%> (-75.52%)`	⬇️
torch_geometric/nn/models/dimenet.py	`14.51% <0.00%> (-53.00%)`	⬇️
torch_geometric/nn/glob/glob.py	`60.52% <0.00%> (-26.32%)`	⬇️
torch_geometric/nn/conv/utils/typing.py	`81.25% <0.00%> (-17.50%)`	⬇️
torch_geometric/profile/profile.py	`32.94% <0.00%> (-15.30%)`	⬇️
torch_geometric/nn/inits.py	`67.85% <0.00%> (-7.15%)`	⬇️
torch_geometric/nn/resolver.py	`88.00% <0.00%> (-6.00%)`	⬇️
torch_geometric/transforms/add_self_loops.py	`94.44% <0.00%> (-5.56%)`	⬇️
... and 11 more

Help us with your feedback. Take ten seconds to tell us how you rate us.

Padarn · 2022-07-22T11:18:18Z

@rusty1s I've added a comment linking to this issue and added details in the description here. I think it is most useful for people looking at this later to be able to come back to this issue if they have any concerns.

torch_geometric/nn/fx.py

rusty1s

Would like to remove the warning before merging - otherwise LGTM. Thanks!

torch_geometric/nn/fx.py

Padarn · 2022-07-23T06:13:23Z

Will do. Thanks!

potential fix

b398ebd

rusty1s reviewed Jul 21, 2022

View reviewed changes

rusty1s assigned Padarn Jul 21, 2022

rusty1s added bug 0 - Priority P0 nn labels Jul 21, 2022

nasty fix

cd67ba4

Padarn changed the title ~~[Discuss] Fix fx problem for Aggregation~~ Fix fx problem for Aggregation Jul 22, 2022

Padarn added 2 commits July 22, 2022 19:16

description

fcfc47b

Merge branch 'master' into padarn/fx-fix

fdaa5a2

change import

295b519

Padarn commented Jul 22, 2022

View reviewed changes

torch_geometric/nn/fx.py Outdated Show resolved Hide resolved

change import

2157bc1

rusty1s approved these changes Jul 23, 2022

View reviewed changes

torch_geometric/nn/fx.py Outdated Show resolved Hide resolved

torch_geometric/nn/fx.py Show resolved Hide resolved

remove warning and move import

6827cd9

Padarn enabled auto-merge (squash) July 23, 2022 06:47

Padarn merged commit 06dbf5b into pyg-team:master Jul 23, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix `fx` problem for `Aggregation` #5021

Fix `fx` problem for `Aggregation` #5021

Padarn commented Jul 21, 2022 •

edited

Loading

rusty1s Jul 21, 2022

rusty1s Jul 21, 2022

Padarn Jul 21, 2022

Padarn Jul 21, 2022

rusty1s Jul 21, 2022

Padarn Jul 21, 2022

rusty1s Jul 21, 2022

Padarn Jul 21, 2022 •

edited

Loading

rusty1s Jul 21, 2022

Padarn Jul 21, 2022

rusty1s Jul 21, 2022

Padarn Jul 22, 2022

rusty1s Jul 22, 2022

Padarn Jul 22, 2022

rusty1s Jul 22, 2022

Padarn Jul 22, 2022

codecov bot commented Jul 22, 2022 •

edited

Loading

Padarn commented Jul 22, 2022

rusty1s left a comment

Padarn commented Jul 23, 2022

Fix fx problem for Aggregation #5021

Fix fx problem for Aggregation #5021

Conversation

Padarn commented Jul 21, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Padarn Jul 21, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Jul 22, 2022 • edited Loading

Codecov Report

Padarn commented Jul 22, 2022

rusty1s left a comment

Choose a reason for hiding this comment

Padarn commented Jul 23, 2022

Fix `fx` problem for `Aggregation` #5021

Fix `fx` problem for `Aggregation` #5021

Padarn commented Jul 21, 2022 •

edited

Loading

Padarn Jul 21, 2022 •

edited

Loading

codecov bot commented Jul 22, 2022 •

edited

Loading