fix CI issues with AOTAutograd

Author: Chillee

functorch/_src/compile_utils.py

@@ -4,6 +4,14 @@
 from torch.utils._pytree import tree_flatten
 
 aten = torch.ops.aten
+
+
+def get_aten_target(node):
+    if hasattr(node.target, 'overloadpacket'):
+        return node.target.overloadpacket
+    return node.target
+
+
 rand_ops = [aten.dropout, aten._fused_dropout, aten._standard_gamma,
             aten.bernoulli, aten.multinomial, aten.native_dropout,
             aten.normal, aten.poisson, aten.binomial, aten.rrelu,
@@ -19,7 +27,7 @@ def fx_graph_cse(fx_g: torch.fx.graph.Graph):
     for n in fx_g.nodes:
         # The placeholder, output, and get_attr nodes are copied to the new grpah without change
         # do not CSE away random operations
-        if n.op == 'placeholder' or n.op == 'output' or n.op == 'get_attr' or n.target in rand_ops:
+        if n.op == 'placeholder' or n.op == 'output' or n.op == 'get_attr' or get_aten_target(n) in rand_ops:
             new_node = new_graph.node_copy(n, lambda x: env[x])
             env[n] = new_node
         else:  # n.op == 'call_function', should never see n.op == 'call_module' or 'call_method'

functorch/_src/partitioners.py

@@ -7,7 +7,7 @@
 import os
 from torch.fx.passes import graph_drawer
 from typing import Tuple
-from .compile_utils import fx_graph_cse, strip_overloads
+from .compile_utils import fx_graph_cse, get_aten_target
 
 
 class InvalidNodeBase(object):
@@ -227,7 +227,6 @@ def min_cut_rematerialization_partition(
     except ImportError:
         raise RuntimeError("Need networkx installed to perform smart recomputation heuristics")
 
-    strip_overloads(joint_module)
     joint_module.graph.eliminate_dead_code()
     joint_module.recompile()
     fx_g = joint_module.graph
@@ -299,22 +298,22 @@ def classify_nodes(joint_module):
 
     def ban_recomputation(node):
         if AGGRESSIVE_RECOMPUTATION:
-            return (node.op == 'call_function' and node.target in unrecomputable_ops)
+            return (node.op == 'call_function' and get_aten_target(node) in unrecomputable_ops)
         else:
             if node.op != 'call_function':
                 return False
-            if node.target not in recomputable_ops:
+            if get_aten_target(node) not in recomputable_ops:
                 return True
             # If the output of the reduction is 4x smaller (arbitrary choice),
             # then we don't allow recomputation.
-            if node.target in reduction_ops:
+            if get_aten_target(node) in reduction_ops:
                 input_tensors_size = sum(_size_of(i.meta['tensor_meta']) for i in node.args if isinstance(i, fx.Node))
                 output_size = _size_of(node.meta['tensor_meta'])
                 return (output_size * 4 < input_tensors_size)
             return False
 
     def is_fusible(a, b):
-        return a.target in fusible_ops and b.target in fusible_ops
+        return get_aten_target(a) in fusible_ops and get_aten_target(b) in fusible_ops
 
     def is_materialized(node):
         if node.op == 'placeholder':

test/test_memory_efficient_fusion.py

@@ -239,8 +239,6 @@ def f(x):
         t = torch.randn(2, 2)
         check(f, t, 0)
 
-    # https://github.com/pytorch/functorch/issues/913
-    @unittest.expectedFailure
     def test_rand_like(self):
         def f(x):
             a = torch.rand_like(x)
@@ -249,17 +247,13 @@ def f(x):
         t = torch.randn(2, 2)
         check(f, t, 0, check_val=False)
 
-    # https://github.com/pytorch/functorch/issues/913
-    @unittest.expectedFailure
     def test_rand_n(self):
         def f(x):
-            g_cpu = torch.Generator()
-            g_cpu.manual_seed(2147483647)
-            a = torch.randn(4, generator=g_cpu)
-            b = torch.randn(4, generator=g_cpu)
+            a = torch.randn(4)
+            b = torch.randn(4)
             return a + b
         t = torch.randn(2, 2)
-        check(f, t, 0)
+        check(f, t, 0, check_val=False)
 
 
 class ReduceTestCase(TestCase):

test/test_minifier.py

@@ -7,7 +7,6 @@
 
 class TestMinifier(TestCase):
     # https://github.com/pytorch/functorch/issues/913
-    @unittest.expectedFailure
     def test_has_mul_minifier(self):
         def failing_f(x, y):
             y = y / 3
@@ -18,7 +17,7 @@ def failing_f(x, y):
         failing_f = make_fx(failing_f)(*inps)
 
         def pass_checker(fx_g, inps):
-            return (torch.ops.aten.mul in set([i.target for i in fx_g.graph.nodes]))
+            return (torch.ops.aten.mul.Tensor in set([i.target for i in fx_g.graph.nodes]))
 
         min_f, inps = minifier(failing_f, inps, pass_checker)
         assert len(min_f.graph.nodes) == 4

test/test_pythonkey.py

@@ -501,7 +501,7 @@ def f(x):
         self.assertEqual(get_num_ins_outs(fw_graph), (1, 3))
 
         ins, outs = get_ins_outs(fw_graph)
-        self.assertEqual(outs[1].target, torch.ops.aten.mm)
+        self.assertEqual(outs[1].target, torch.ops.aten.mm.default)
 
 
 class TestContiguous(TestCase):