Checking hasattr _c for script (#105)

Author: XiaobingSuper

intel_pytorch_extension_py/ops/jit.py

@@ -14,7 +14,7 @@ def script_(obj, optimize=None, _frames_up=0, _rcb=None):
     jit_m = orig_script(obj, optimize=optimize, _frames_up=_frames_up+1, _rcb=_rcb)
     torch.jit.script = script_
 
-    if core.get_jit_opt() and isinstance(jit_m, torch._C.ScriptModule):
+    if core.get_jit_opt() and hasattr(jit_m, '_c'):
         # Disable mix precision in model fusion, since mixed precision cannot
         # bring any benefits for inference, but will lead to loss of accuracy
         orig_mixed_type = ipex.get_auto_mix_precision()
@@ -31,7 +31,7 @@ def trace_(func, example_inputs, *args, **kwargs):
     ipex.enable_auto_mix_precision(None)
     jit_m = orig_trace(func, example_inputs, *args, **kwargs)
 
-    if core.get_jit_opt() and isinstance(jit_m, torch._C.ScriptModule):
+    if core.get_jit_opt() and hasattr(jit_m, '_c'):
         jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
     ipex.enable_auto_mix_precision(orig_mixed_type)
     return jit_m

tests/cpu/test_jit.py

@@ -192,7 +192,7 @@ def _test_output(self, model, x, kind_in_graph=None, kind_not_in_graph=None):
         if kind_in_graph is not None:
             self.assertTrue(any(n.kind() == kind_in_graph for n in script_graph.nodes()))
             self.assertTrue(any(n.kind() == kind_in_graph for n in trace_graph.nodes()))
-        
+
         # check if certain node does not exist in the graph
         if kind_not_in_graph is not None:
             self.assertTrue(all(n.kind() != kind_not_in_graph for n in script_graph.nodes()))
@@ -205,12 +205,12 @@ def _test_output_bf16(self, model, x, kind_in_graph=None, kind_not_in_graph=None
         core.enable_auto_dnnl()
         core.enable_jit_opt()
         core.enable_mix_bf16_fp32()
-        
+
         model = model.to(ipex.DEVICE).eval()
         x = x.to(ipex.DEVICE)
         x2 = x.clone()
         x3 = x.clone()
-        
+
         script_fused_model = torch.jit.script(copy.deepcopy(model))
         trace_fused_model = torch.jit.trace(copy.deepcopy(model), x3)
 
@@ -224,24 +224,24 @@ def _test_output_bf16(self, model, x, kind_in_graph=None, kind_not_in_graph=None
             # bf 16, jit trace path
             trace_graph = trace_fused_model.graph_for(x3)
             fused_tresult = trace_fused_model(x3)
-        
-        # disable mix_bf16_fp32 when the calculation is done 
+
+        # disable mix_bf16_fp32 when the calculation is done
         # to avoid affecting other scripts
         core.disable_mix_bf16_fp32()
-        
+
         self.assertEqual(fused_sresult, result, prec=prec)
         self.assertEqual(fused_tresult, result, prec=prec)
 
         # check if the fused node exists in the graph
         if kind_in_graph is not None:
             self.assertTrue(any(n.kind() == kind_in_graph for n in script_graph.nodes()))
             self.assertTrue(any(n.kind() == kind_in_graph for n in trace_graph.nodes()))
-        
+
         # check if certain node does not exist in the graph
         if kind_not_in_graph is not None:
             self.assertTrue(all(n.kind() != kind_not_in_graph for n in script_graph.nodes()))
             self.assertTrue(all(n.kind() != kind_not_in_graph for n in trace_graph.nodes()))
-    
+
 
     def test_output_conv_bn_2d(self):
         self._test_output(
@@ -364,6 +364,23 @@ def test_output_linear_relu(self):
             kind_in_graph="ipex::linear_relu")
 
 
+    def test_jit_function(self):
+        # test hool trace and script can works for function
+        def fn(input, weight, bias):
+            return F.linear(input, weight, bias)
+
+        input = torch.randn(2, 4)
+        weight = torch.randn(5, 4)
+        bias = torch.randn(5)
+        result = fn(input, weight, bias)
+
+        scripted_fn = torch.jit.script(fn)
+        traced_fn = torch.jit.trace(fn, (input, weight, bias))
+
+        self.assertEqual(scripted_fn(input, weight, bias), result)
+        self.assertEqual(traced_fn(input, weight, bias), result)
+
+
 if __name__ == '__main__':
     torch.manual_seed(2020)
     core.enable_auto_dnnl()