Only require python print on certain namespaces (#18846)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18846
ghimport-source-id: b211e15d24c88fdc32d79222d9fce2fa9c291541

Differential Revision: D14901485

Pulled By: bwasti

fbshipit-source-id: 4b02a66d325ba5391d1f838055aea13b5e4f6485

diff --git a/torch/csrc/jit/passes/python_print.cpp b/torch/csrc/jit/passes/python_print.cpp
index f7bfd99..40dc2d3 100644 (file)
--- a/torch/csrc/jit/passes/python_print.cpp
+++ b/torch/csrc/jit/passes/python_print.cpp
@@ -1263,7 +1263,16 @@ TORCH_API bool printerHasSpecialCaseFor(Symbol sym) {
 
   };
 
-  return handled.count(sym) || unneeded.count(sym);
+  // These namespaces are required to have Python printers unless
+  // otherwise noted in unneeded.
+  const static std::unordered_set<Symbol> required_namespaces = {
+      c10::namespaces::prim,
+      c10::namespaces::aten,
+      c10::namespaces::onnx,
+  };
+
+  return handled.count(sym) || unneeded.count(sym) ||
+      !required_namespaces.count(sym.ns());
 }
 
 } // namespace jit