Set __file__ for torch.ops (#21888)

Ailing Zhang · facebook-github-bot · commit f363a33e1033 · 2019-06-18T08:46:23.000-07:00
Summary: Fixes #19351 meta-pytorch/lockdown#93 Pull Request resolved: #21888 Differential Revision: D15871142 Pulled By: ailzhang fbshipit-source-id: 339e9d493e2e13f09e118814bdd1d7a5942804b8
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -2829,6 +2829,13 @@ def test_pytorch_jit_env_off(self):
         except subprocess.CalledProcessError as e:
             raise RuntimeError("Could not 'import torch' with PYTORCH_JIT=0")
 
+    def test_print_op_module(self):
+        # Issue #19351: python2 and python3 go through different paths.
+        # python2 returns '<module 'torch.ops' (built-in)>'
+        # python3 uses __file__ and return
+        # '<module 'torch.ops' from '/scratch/ailzhang/pytorch/torch/_ops.py'>'
+        s = str(torch.ops)
+        self.assertRegex(s, r'ops')
 
 def execWrapper(code, glob, loc):
     if PY2:
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -2,6 +2,7 @@
 
 import contextlib
 import ctypes
+import os
 import sys
 import types
 
@@ -66,6 +67,8 @@ def __getattr__(self, op_name):
 
 
 class _Ops(types.ModuleType):
+    __file__ = os.path.join(os.path.dirname(__file__), '_ops.py')
+
     def __init__(self):
         super(_Ops, self).__init__('torch.ops')
         self.loaded_libraries = set()
