| 81 | // --------------------------------------------------------------------------- |
| 82 | |
| 83 | static decltype(torch::OperatorRegistry::instance() |
| 84 | .find_operator("") |
| 85 | ->implementations.end()) |
| 86 | pick_impl(torch::OperatorRegistration* op) { |
| 87 | using DK = c10::DispatchKey; |
| 88 | const std::vector<DK> preferred_keys = { |
| 89 | DK::CPU, DK::BackendSelect, DK::CatchAll}; |
| 90 | auto chosen = op->implementations.end(); |
| 91 | for (const auto& key : preferred_keys) { |
| 92 | chosen = op->implementations.find(key); |
| 93 | if (chosen != op->implementations.end()) break; |
| 94 | } |
| 95 | // Mirror the production rule: allow exactly-one-impl, reject ambiguous. |
| 96 | if (chosen == op->implementations.end() && op->implementations.size() == 1) { |
| 97 | chosen = op->implementations.begin(); |
| 98 | } |
| 99 | return chosen; |
| 100 | } |
| 101 | |
| 102 | // --------------------------------------------------------------------------- |
| 103 | // Tests |