Add TorchExpm1Visitor (#78)

Follow-up for #77

Author: kit1980

tests/fixtures/misc/checker/expm1.py

@@ -0,0 +1,12 @@
+import torch
+a = torch.randn(5)
+b = torch.exp(a) - 1
+c = torch.exp(a) - 1.0
+
+ret = (torch.exp(a) - 1) * torch.exp(2 * b)
+
+# False negative: can not detect currently
+x = a.exp() - 1
+
+# False negative: should be rare and would complicate implementation
+x = -1 + torch.exp(a)

tests/fixtures/misc/checker/expm1.txt

@@ -0,0 +1,3 @@
+3:5 TOR107 Use `torch.special.expm1(x)` instead of `torch.exp(x) - 1`. It is more accurate for small values of `x`.
+4:5 TOR107 Use `torch.special.expm1(x)` instead of `torch.exp(x) - 1`. It is more accurate for small values of `x`.
+6:7 TOR107 Use `torch.special.expm1(x)` instead of `torch.exp(x) - 1`. It is more accurate for small values of `x`.

tests/test_torchfix.py

@@ -37,7 +37,15 @@ def pytest_generate_tests(metafunc):
             ("TOR102,TOR101", {"TOR102", "TOR101"}),
             (
                 "TOR1,TOR102",
-                {"TOR102", "TOR101", "TOR103", "TOR104", "TOR105", "TOR106"},
+                {
+                    "TOR101",
+                    "TOR102",
+                    "TOR103",
+                    "TOR104",
+                    "TOR105",
+                    "TOR106",
+                    "TOR107",
+                },
             ),
             (None, set(GET_ALL_ERROR_CODES()) - exclude_set),
         ]

torchfix/torchfix.py

@@ -9,6 +9,7 @@
 
 from .visitors import (
     TorchDeprecatedSymbolsVisitor,
+    TorchExpm1Visitor,
     TorchLog1pVisitor,
     TorchNonPublicAliasVisitor,
     TorchReentrantCheckpointVisitor,
@@ -29,6 +30,7 @@
 
 ALL_VISITOR_CLS = [
     TorchDeprecatedSymbolsVisitor,
+    TorchExpm1Visitor,
     TorchLog1pVisitor,
     TorchNonPublicAliasVisitor,
     TorchRequireGradVisitor,

torchfix/visitors/__init__.py

@@ -1,9 +1,10 @@
 from .deprecated_symbols import TorchDeprecatedSymbolsVisitor
 from .internal import TorchScopedLibraryVisitor
 from .misc import (
+    TorchExpm1Visitor,
+    TorchLog1pVisitor,
     TorchReentrantCheckpointVisitor,
     TorchRequireGradVisitor,
-    TorchLog1pVisitor,
 )
 from .nonpublic import TorchNonPublicAliasVisitor
 from .performance import TorchSynchronizedDataLoaderVisitor
@@ -16,6 +17,7 @@
 
 __all__ = [
     "TorchDeprecatedSymbolsVisitor",
+    "TorchExpm1Visitor",
     "TorchLog1pVisitor",
     "TorchNonPublicAliasVisitor",
     "TorchReentrantCheckpointVisitor",

torchfix/visitors/misc/__init__.py

@@ -121,3 +121,36 @@ def visit_Call(self, node):
                     message=self.ERRORS[0].message(),
                     replacement=None,
                 )
+
+
+class TorchExpm1Visitor(TorchVisitor):
+    """
+    Suggest using `torch.special.expm1(x)` instead of `torch.exp(x) - 1`.
+    """
+
+    ERRORS = [
+        TorchError(
+            "TOR107",
+            (
+                "Use `torch.special.expm1(x)` instead of `torch.exp(x) - 1`. "
+                "It is more accurate for small values of `x`."
+            ),
+        )
+    ]
+
+    def visit_BinaryOperation(self, node):
+        if m.matches(
+            node,
+            m.BinaryOperation(
+                left=m.Call(),
+                operator=m.Subtract(),
+                right=m.Integer(value="1") | m.Float(value="1.0"),
+            ),
+        ):
+            if self.get_qualified_name_for_call(node.left) == "torch.exp":
+                self.add_violation(
+                    node,
+                    error_code=self.ERRORS[0].error_code,
+                    message=self.ERRORS[0].message(),
+                    replacement=None,
+                )