Add test to torch.jit.export_opnames, make the _C function private

Summary: Pull Request resolved: pytorch#31446 Test Plan: Imported from OSS Differential Revision: D19172851 Pulled By: iseeyuan fbshipit-source-id: f06d8766ed73c9abe4ebf41c402ee64880d745be

Author: Martin Yuan

docs/source/torch.rst

@@ -377,4 +377,3 @@ Utilities
 .. autofunction:: result_type
 .. autofunction:: can_cast
 .. autofunction:: promote_types
-.. autofunction:: export_opnames

test/test_jit.py

@@ -3693,6 +3693,38 @@ def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh):
  self.assertTrue(type(block.paramNode()) == torch._C.Node)
  self.assertTrue(tested_blocks)
 
+ def test_export_opnames(self):
+ class Foo(torch.jit.ScriptModule):
+ def __init__(self):
+ super(Foo, self).__init__()
+
+ def one(self, x, y):
+ # type: (Tensor, Tensor) -> Tensor
+ return x + y
+
+ def two(self, x):
+ # type: (Tensor) -> Tensor
+ return 2 * x
+
+ @torch.jit.script_method
+ def forward(self, x):
+ # type: (Tensor) -> Tensor
+ return self.one(self.two(x), x)
+
+ class Bar(torch.jit.ScriptModule):
+ def __init__(self):
+ super(Bar, self).__init__()
+ self.sub = Foo()
+
+ def forward(self, x):
+ # type: (Tensor) -> Tensor
+ return self.sub.forward(x)
+
+ bar = Bar()
+ ops = torch.jit.export_opnames(bar)
+ expected = ['aten::add.Tensor', 'aten::mul.Scalar', 'prim::Constant']
+ self.assertEqual(ops, expected)
+
  def test_pytorch_jit_env_off(self):
  import subprocess
  env = os.environ.copy()
@@ -3735,7 +3767,6 @@ def forward(self):
  model_loaded = torch.jit.load(buffer)
  self.assertEqual(model_loaded(), model())
 
-
 class TestFrontend(JitTestCase):
 
  def test_instancing_error(self):

torch/csrc/jit/script/init.cpp

@@ -1175,7 +1175,7 @@ void initJitScriptBindings(PyObject* module) {
  return Module(get_python_cu(), type);
  });
 
- m.def("export_opnames",
+ m.def("_export_opnames",
  [](script::Module& sm) {return debugMakeList(torch::jit::export_opnames(sm));});
 
  py::class_<ConcreteModuleTypeBuilder, std::shared_ptr<ConcreteModuleTypeBuilder>>(

torch/jit/__init__.py

@@ -240,7 +240,7 @@ def export_opnames(m):
  r"""
  Returns a list of operator names of a script module and its submodules
  """
- return torch._C.export_opnames(m._c)
+ return torch._C._export_opnames(m._c)
 
 def _get_trace_graph(f, args=(), kwargs=None, _force_outplace=False,
  return_inputs=False, _return_inputs_states=False):