diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in index 0218c9fbdf4..70248d13252 100644 --- a/torch/_C/__init__.pyi.in +++ b/torch/_C/__init__.pyi.in @@ -1052,6 +1052,7 @@ def _dispatch_has_kernel_for_dispatch_key(name: str, dispatch: _dispatchkey) -> def _dispatch_has_kernel_for_any_dispatch_key(name: str, dispatch_key_set: DispatchKeySet) -> _bool: ... def _dispatch_has_computed_kernel_for_dispatch_key(name: str, dispatch: _dispatchkey) -> _bool: ... def _dispatch_find_dangling_impls() -> List[str]: ... +def _dispatch_get_all_op_names() -> List[str]: ... def _dispatch_tls_set_dispatch_key_excluded(dispatch: _dispatchkey, val: _bool) -> None: ... def _dispatch_tls_is_dispatch_key_excluded(dispatch: _dispatchkey) -> _bool: ... def _dispatch_isTensorSubclassLike(tensor: Tensor) -> _bool: ... diff --git a/torch/csrc/utils/python_dispatch.cpp b/torch/csrc/utils/python_dispatch.cpp index 96464683d38..4ce1d3a1c5d 100644 --- a/torch/csrc/utils/python_dispatch.cpp +++ b/torch/csrc/utils/python_dispatch.cpp @@ -340,6 +340,23 @@ void initDispatchBindings(PyObject* module) { return states; }); + m.def("_dispatch_get_all_op_names", []() -> std::vector { + auto op_names = c10::Dispatcher::singleton().getAllOpNames(); + + std::vector names; + names.reserve(op_names.size()); + for (auto& op : op_names) { + std::stringstream ss; + ss << op.name; + if (!op.overload_name.empty()) { + ss << "." << op.overload_name; + } + names.push_back(ss.str()); + } + + return names; + }); + m.def( "_dispatch_tls_set_dispatch_key_excluded", [](c10::DispatchKey dispatch_key, bool desired_state) {