[TorchFix] Deprecated codemods to honor aliases (#4680)

By introducing `torchfix.common.get_module_name(node: cst.Call)` and
using it in deprecated symbols codemods

Fixes pytorch/test-infra#4452

Author: malfet

tests/fixtures/deprecated_symbols/codemod/aliased_import.py

@@ -0,0 +1,10 @@
+import torch
+import torch as foo
+import torch as bar
+import torch as baz
+
+A = torch.arange(9.0).reshape(3, 3)
+A3 = foo.chain_matmul(A, A, A)
+rc1 = bar.cholesky(torch.mm(A3.t(), A3))
+rc2 = baz.qr(torch.mm(A.t(), A))
+

tests/fixtures/deprecated_symbols/codemod/aliased_import.py.out

@@ -0,0 +1,10 @@
+import torch
+import torch as foo
+import torch as bar
+import torch as baz
+
+A = torch.arange(9.0).reshape(3, 3)
+A3 = foo.linalg.multi_dot([A, A, A])
+rc1 = bar.linalg.cholesky(torch.mm(A3.t(), A3))
+rc2 = baz.linalg.qr(torch.mm(A.t(), A))
+

torchfix/common.py

@@ -120,3 +120,11 @@ def on_leave(self, original_node, updated_node):
             return updated_node
 
     return tree.visit(MultiChildReplacementTransformer(replacement_map))
+
+
+def get_module_name(node: cst.Call, default: Optional[str] = None) -> Optional[str]:
+    if not isinstance(node.func, cst.Attribute):
+        return default
+    if not isinstance(node.func.value, cst.Name):
+        return default
+    return node.func.value.value

torchfix/visitors/deprecated_symbols/chain_matmul.py

@@ -1,4 +1,5 @@
 import libcst as cst
+from ...common import get_module_name
 
 
 def call_replacement_chain_matmul(node: cst.Call) -> cst.CSTNode:
@@ -19,7 +20,8 @@ def call_replacement_chain_matmul(node: cst.Call) -> cst.CSTNode:
         replacement_args = [matrices_arg]
     else:
         replacement_args = [matrices_arg, out_arg]
-    replacement = cst.parse_expression("torch.linalg.multi_dot(args)")
+    module_name = get_module_name(node, 'torch')
+    replacement = cst.parse_expression(f"{module_name}.linalg.multi_dot(args)")
     replacement = replacement.with_changes(args=replacement_args)
 
     return replacement

torchfix/visitors/deprecated_symbols/cholesky.py

@@ -1,5 +1,5 @@
 import libcst as cst
-from ...common import TorchVisitor
+from ...common import (TorchVisitor, get_module_name)
 
 
 def call_replacement_cholesky(node: cst.Call) -> cst.CSTNode:
@@ -12,20 +12,20 @@ def call_replacement_cholesky(node: cst.Call) -> cst.CSTNode:
         comma=cst.MaybeSentinel.DEFAULT
     )
     upper_arg = TorchVisitor.get_specific_arg(node, "upper", 1)
+    module_name = get_module_name(node, "torch")
 
     if (
         upper_arg is not None
         and cst.ensure_type(upper_arg.value, cst.Name).value == "True"
     ):
-        replacement = cst.parse_expression("torch.linalg.cholesky(A).mH")
+        replacement = cst.parse_expression(f"{module_name}.linalg.cholesky(A).mH")
         replacement = replacement.with_deep_changes(
             # Ignore type error, see https://github.com/Instagram/LibCST/issues/963
             old_node=cst.ensure_type(replacement.value, cst.Call).args,  # type: ignore
             value=[input_arg],
         )
     else:
-        replacement = cst.parse_expression("torch.linalg.cholesky(A)").with_changes(
-            args=[input_arg]
-        )
+        replacement = cst.parse_expression(f"{module_name}.linalg.cholesky(A)")
+        replacement = replacement.with_changes(args=[input_arg])
 
     return replacement

torchfix/visitors/deprecated_symbols/qr.py

@@ -1,6 +1,6 @@
 import libcst as cst
 from typing import Optional
-from ...common import TorchVisitor
+from ...common import (TorchVisitor, get_module_name)
 
 
 def call_replacement_qr(node: cst.Call) -> Optional[cst.CSTNode]:
@@ -27,7 +27,8 @@ def call_replacement_qr(node: cst.Call) -> Optional[cst.CSTNode]:
             comma=cst.MaybeSentinel.DEFAULT
         )
         replacement_args = [input_arg]
-    replacement = cst.parse_expression("torch.linalg.qr(args)")
+    module_name = get_module_name(node, "torch")
+    replacement = cst.parse_expression(f"{module_name}.linalg.qr(args)")
     replacement = replacement.with_changes(args=replacement_args)
 
     return replacement