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

Added "dump ops" API to return ops instead of print (#78995)

Browse Source 
Useful to use for grabbing info instead of the hacky "redirect C++ output" stuff I currently do.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/78995
Approved by: https://github.com/ezyang
This commit is contained in:
Horace He 2022-06-07 05:19:05 +00:00 committed by PyTorch MergeBot
parent 69778ee4eb
commit bbbfbbeddc
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -248,6 +248,17 @@ void initDispatchBindings(PyObject* module) {
std::cout << op << std::endl;
}
}, py::arg("dispatch_key") = static_cast<const char*>(""));
m.def("_dispatch_get_registrations_for_dispatch_key", [](const char* dispatch_key = "") {
auto k = std::string(dispatch_key) == "" ? c10::nullopt : c10::make_optional(c10::parseDispatchKey(dispatch_key));
auto op_names = c10::Dispatcher::singleton().getRegistrationsForDispatchKey(k);
std::vector<std::string> names;
names.reserve(op_names.size());
for (auto& op : op_names) {
names.push_back(op.name + (op.overload_name == "" ? "" : "." + op.overload_name));
}
return names;
}, py::arg("dispatch_key") = static_cast<const char*>(""));
}
}}} // namespace torch::impl::dispatch