OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Python binding for dispatcher getAllOpNames (#87422)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87422
Approved by: https://github.com/bdhirsh
This commit is contained in:
Sherlock Huang 2022-10-21 00:46:34 +00:00 committed by PyTorch MergeBot
parent 7caeac1718
commit ab901b4817
2 changed files with 18 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
torch/_C/__init__.pyi.in
View File

@ -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: ...

17
torch/csrc/utils/python_dispatch.cpp
View File

@ -340,6 +340,23 @@ void initDispatchBindings(PyObject* module) {
return states;
});
m.def("_dispatch_get_all_op_names", []() -> std::vector<std::string> {
auto op_names = c10::Dispatcher::singleton().getAllOpNames();
std::vector<std::string> 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) {