updates per review

Mike Ruberry · Mike Ruberry · commit 38afbbab40b2 · 2020-08-10T23:24:54.000-07:00
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -56,7 +56,7 @@
     execWrapper, _inline_everything, _tmp_donotuse_dont_inline_everything, \
     RUN_CUDA
 from torch.testing._internal.jit_metaprogramming_utils import create_script_fn, nn_functional_tests, get_script_args, \
-    EXCLUDE_SCRIPT, additional_module_tests, EXCLUDE_SCRIPT_MODULES, \
+    get_call, script_template, EXCLUDE_SCRIPT, additional_module_tests, EXCLUDE_SCRIPT_MODULES, \
     get_nn_module_name_from_kwargs, script_method_template, create_traced_fn
 
 from torch.testing._internal.common_nn import module_tests, new_module_tests, criterion_tests
@@ -15165,6 +15165,14 @@ def parameter_script(x: torch.nn.Parameter):
     'test_nn_max_pool1d_with_indices',
 }
 
+def check_alias_annotation(method_name, args, kwargs):
+    formals, tensors, actuals = get_script_args(args)
+    call = get_call(method_name, 'method', actuals, kwargs)
+    script = script_template.format(', '.join(formals), call)
+    CU = torch.jit.CompilationUnit(script)
+    torch._C._jit_check_alias_annotation(CU.the_method.graph, tuple(tensors), method_name)
+
+
 def check_output_types(self, func, ref_outputs, args, kwargs):
     graph = getattr(func, 'last_graph', None)
     types = [o.type() for o in graph.outputs()]
@@ -15427,6 +15435,10 @@ def fn(*inputs, **kwargs):
                                                     fn, f_args_variable, kwargs_variable,
                                                     check_types=check_types)
 
+                # alias annotation testing
+                if is_inplace and test_name not in EXCLUDE_SCRIPT:
+                    check_alias_annotation(name, (self_variable,) + args_variable, kwargs_variable)
+
             check(name)
             inplace_name = name + '_'
             # can't broadcast inplace to left hand side
diff --git a/test/test_linalg.py b/test/test_linalg.py
@@ -20,12 +20,8 @@ def test_outer(self, device, dtype):
         a = torch.randn(50, device=device, dtype=dtype)
         b = torch.randn(50, device=device, dtype=dtype)
 
-        def _fn(a, b):
-            return torch.outer(a, b)
-
         ops = (torch.ger, torch.Tensor.ger,
-               torch.outer, torch.Tensor.outer,
-               torch.jit.script(_fn))
+               torch.outer, torch.Tensor.outer)
 
         expected = np.outer(a.cpu().numpy(), b.cpu().numpy())
         for op in ops:
@@ -44,11 +40,9 @@ def test_det(self, device, dtype):
             torch.randn((3, 52, 52), device=device, dtype=dtype),
             torch.randn((4, 2, 26, 26), device=device, dtype=dtype))
 
-        def _fn(t):
-            return torch.linalg.det(t)
 
         ops = (torch.det, torch.Tensor.det,
-               torch.linalg.det, torch.jit.script(_fn))
+               torch.linalg.det)
         for t in tensors:
             expected = np.linalg.det(t.cpu().numpy())
             for op in ops:
@@ -57,12 +51,9 @@ def _fn(t):
 
         # NOTE: det requires a 2D+ tensor
         t = torch.randn(1, device=device, dtype=dtype)
-        for op in ops:
-            try:
-                op(t)
-            except Exception as e:
-                continue
-            self.assertTrue(False, msg="Failed to throw error on 1D tensor!")
+        with self.assertRaises(IndexError):
+            op(t)
+
 
 instantiate_device_type_tests(TestLinalg, globals())
 
diff --git a/test/test_op_normalization.py b/test/test_op_normalization.py
@@ -58,9 +58,11 @@ def _test(self, device, info=info):
             op = info.alias_op
             is_inplace = info.alias_name.endswith('_')
 
-            # NOTE: this workaround necessary to satisfy the JIT, which
-            #   does allow Python aliasing or direct calling of torch.Tensor
-            #   methods when scripting
+            # Checks that scripting converts aliases
+            # NOTE: the code to test scripting must be generated since
+            #   scripting does not support splatting args or directly
+            #   calling torch.Tensor methods. The following
+            #   splats args after the first tensor by inlining them as constants.
             if is_inplace:
                 fn_template = '''
                     def _fn(t):
@@ -69,9 +71,6 @@ def _fn(t):
                 arg_string = ', '.join((str(arg) for arg in info.args))
                 script = fn_template.format(alias_name=info.alias_name, args=arg_string)
             else:
-                # Checks that scripting converts aliases
-                # NOTE: scripting doesn't support splatting args, so this
-                #  generates the script with the args already splatted (as constants)
                 fn_template = '''
                     def _fn(t):
                         return op(t{args})
