Don't suggest logsumexp if sum's dim is None (#91)

See pytorch/pytorch#144339

Author: kit1980

tests/fixtures/misc/checker/logsumexp.py

@@ -1,14 +1,21 @@
 import torch
-a = torch.randn(5)
-b = torch.randn(5)
+
+x = torch.randn(5)
 
 # logsumexp
 y = torch.log(torch.sum(torch.exp(x), 1, keepdim=True))
+y = torch.log(torch.sum(torch.exp(x), dim=1, keepdim=True))
 y = torch.log(torch.sum(torch.exp(2.5 + x), 1))
+y = torch.log(torch.sum(torch.exp(2.5 + x), dim=1))
 
 # not logsumexp
 y = torch.log(torch.sum(torch.exp(x), 1, keepdim=True) + 2.5)
 y = torch.log(torch.sum(torch.exp(x) + 2.5, 1))
 y = torch.log(2 + x)
 y = torch.sum(torch.log(torch.exp(x)), 1)
 y = torch.exp(torch.sum(torch.log(x), 1, keepdim=True))
+
+# not logsumexp because of https://github.com/pytorch/pytorch/issues/144339
+y = torch.log(torch.sum(torch.exp(x), None, keepdim=True))
+y = torch.log(torch.sum(torch.exp(x), dim=None, keepdim=True))
+y = torch.log(torch.sum(torch.exp(x), keepdim=True))

tests/fixtures/misc/checker/logsumexp.txt

@@ -1,2 +1,4 @@
 6:5 TOR108 Use numerically stabilized `torch.logsumexp`.
 7:5 TOR108 Use numerically stabilized `torch.logsumexp`.
+8:5 TOR108 Use numerically stabilized `torch.logsumexp`.
+9:5 TOR108 Use numerically stabilized `torch.logsumexp`.

torchfix/visitors/misc/__init__.py

@@ -184,9 +184,20 @@ def visit_Call(self, node):
                         )
                         == "torch.exp"
                     ):
-                        self.add_violation(
-                            node,
-                            error_code=self.ERRORS[0].error_code,
-                            message=self.ERRORS[0].message(),
-                            replacement=None,
+
+                        # if `dim` is not provided or None for sum, skip:
+                        # https://github.com/pytorch/pytorch/issues/144339
+                        dim_arg = self.get_specific_arg(
+                            node.args[0].value, arg_name="dim", arg_pos=1
                         )
+                        if dim_arg is not None:
+                            if not (
+                                isinstance(dim_arg.value, cst.Name)
+                                and dim_arg.value.value == "None"
+                            ):
+                                self.add_violation(
+                                    node,
+                                    error_code=self.ERRORS[0].error_code,
+                                    message=self.ERRORS[0].message(),
+                                    replacement=None,
+                                )