Add TorchLog1pVisitor

tests/fixtures/misc/checker/log1p.py

@@ -0,0 +1,6 @@
+import torch
+a = torch.randn(5)
+b = torch.log(1 + a)
+c = torch.log(a + 1)
+b = torch.log(1.0 + a)
+c = torch.log(a + 1.0)

tests/fixtures/misc/checker/log1p.txt

@@ -0,0 +1,4 @@
+3:5 TOR106 Use `torch.log1p(x)` instead of `torch.log(1 + x)`. It is more accurate for small values of `x`.
+4:5 TOR106 Use `torch.log1p(x)` instead of `torch.log(1 + x)`. It is more accurate for small values of `x`.
+5:5 TOR106 Use `torch.log1p(x)` instead of `torch.log(1 + x)`. It is more accurate for small values of `x`.
+6:5 TOR106 Use `torch.log1p(x)` instead of `torch.log(1 + x)`. It is more accurate for small values of `x`.

tests/test_torchfix.py

@@ -35,7 +35,10 @@ def pytest_generate_tests(metafunc):
             ("ALL,TOR102", GET_ALL_ERROR_CODES()),
             ("TOR102", {"TOR102"}),
             ("TOR102,TOR101", {"TOR102", "TOR101"}),
-            ("TOR1,TOR102", {"TOR102", "TOR101", "TOR103", "TOR104", "TOR105"}),
+            (
+                "TOR1,TOR102",
+                {"TOR102", "TOR101", "TOR103", "TOR104", "TOR105", "TOR106"},
+            ),
             (None, set(GET_ALL_ERROR_CODES()) - exclude_set),
         ]
         metafunc.parametrize("case,expected", cases, ids=[case for case, _ in cases])

torchfix/torchfix.py

@@ -9,6 +9,7 @@
 
 from .visitors import (
     TorchDeprecatedSymbolsVisitor,
+    TorchLog1pVisitor,
     TorchNonPublicAliasVisitor,
     TorchReentrantCheckpointVisitor,
     TorchRequireGradVisitor,
@@ -28,15 +29,16 @@
 
 ALL_VISITOR_CLS = [
     TorchDeprecatedSymbolsVisitor,
+    TorchLog1pVisitor,
+    TorchNonPublicAliasVisitor,
     TorchRequireGradVisitor,
+    TorchReentrantCheckpointVisitor,
     TorchScopedLibraryVisitor,
     TorchSynchronizedDataLoaderVisitor,
+    TorchUnsafeLoadVisitor,
     TorchVisionDeprecatedPretrainedVisitor,
     TorchVisionDeprecatedToTensorVisitor,
     TorchVisionSingletonImportVisitor,
-    TorchUnsafeLoadVisitor,
-    TorchReentrantCheckpointVisitor,
-    TorchNonPublicAliasVisitor,
 ]
 
 

torchfix/visitors/__init__.py

@@ -1,6 +1,10 @@
 from .deprecated_symbols import TorchDeprecatedSymbolsVisitor
 from .internal import TorchScopedLibraryVisitor
-from .misc import TorchReentrantCheckpointVisitor, TorchRequireGradVisitor
+from .misc import (
+    TorchReentrantCheckpointVisitor,
+    TorchRequireGradVisitor,
+    TorchLog1pVisitor,
+)
 from .nonpublic import TorchNonPublicAliasVisitor
 from .performance import TorchSynchronizedDataLoaderVisitor
 from .security import TorchUnsafeLoadVisitor
@@ -12,13 +16,14 @@
 
 __all__ = [
     "TorchDeprecatedSymbolsVisitor",
+    "TorchLog1pVisitor",
+    "TorchNonPublicAliasVisitor",
+    "TorchReentrantCheckpointVisitor",
     "TorchRequireGradVisitor",
     "TorchScopedLibraryVisitor",
     "TorchSynchronizedDataLoaderVisitor",
+    "TorchUnsafeLoadVisitor",
     "TorchVisionDeprecatedPretrainedVisitor",
     "TorchVisionDeprecatedToTensorVisitor",
     "TorchVisionSingletonImportVisitor",
-    "TorchUnsafeLoadVisitor",
-    "TorchReentrantCheckpointVisitor",
-    "TorchNonPublicAliasVisitor",
 ]

torchfix/visitors/misc/__init__.py

@@ -77,3 +77,47 @@ def visit_Call(self, node):
                 message=self.ERRORS[0].message(),
                 replacement=replacement,
             )
+
+
+class TorchLog1pVisitor(TorchVisitor):
+    """
+    Suggest using `torch.log1p(x)` instead of `torch.log(1 + x)`.
+    """
+
+    ERRORS = [
+        TorchError(
+            "TOR106",
+            (
+                "Use `torch.log1p(x)` instead of `torch.log(1 + x)`. "
+                "It is more accurate for small values of `x`."
+            ),
+        )
+    ]
+
+    def visit_Call(self, node):
+        if self.get_qualified_name_for_call(node) == "torch.log":
+
+            if m.matches(
+                node,
+                m.Call(
+                    args=[
+                        m.Arg(
+                            value=m.BinaryOperation(
+                                left=m.Integer(value="1") | m.Float(value="1.0"),
+                                operator=m.Add(),
+                            )
+                            | m.BinaryOperation(
+                                operator=m.Add(),
+                                right=m.Integer(value="1") | m.Float(value="1.0"),
+                            ),
+                        ),
+                    ],
+                ),
+            ):
+
+                self.add_violation(
+                    node,
+                    error_code=self.ERRORS[0].error_code,
+                    message=self.ERRORS[0].message(),
+                    replacement=None,
+                )