if all the inputs are not megengine tensor, return get_default_device() else check whether all input tensors have the same device
| 179 | // if all the inputs are not megengine tensor, return get_default_device() |
| 180 | // else check whether all input tensors have the same device |
| 181 | CompNode _get_device(PyObject* const* args, size_t nargs) { |
| 182 | bool is_tuple = false; |
| 183 | PyObject* tuple = nullptr; |
| 184 | // convert input args to a tuple |
| 185 | if (nargs == 1 && (PyTuple_Check(args[0]) || PyList_Check(args[0]))) { |
| 186 | if (PyList_Check(args[0])) { |
| 187 | tuple = PyList_AsTuple(args[0]); |
| 188 | } else { |
| 189 | tuple = args[0]; |
| 190 | Py_INCREF(tuple); |
| 191 | } |
| 192 | nargs = PyTuple_Size(tuple); |
| 193 | is_tuple = true; |
| 194 | } |
| 195 | bool valid = false; |
| 196 | CompNode cn; |
| 197 | for (size_t i = 0; i < nargs; ++i) { |
| 198 | PyObject* handle = is_tuple ? PyTuple_GetItem(tuple, i) : args[i]; |
| 199 | TensorWrapper* tw = TensorWrapper::try_cast(handle); |
| 200 | |
| 201 | if (tw) { |
| 202 | if (!valid) { |
| 203 | cn = tw->m_tensor->comp_node(); |
| 204 | valid = true; |
| 205 | } else { |
| 206 | CompNode cn1 = tw->m_tensor->comp_node(); |
| 207 | if (cn1 != cn) { |
| 208 | throw py::value_error(ssprintf( |
| 209 | "ambiguous device: %s (from %s) vs %s (from %s)", |
| 210 | cn.to_string().c_str(), cn.to_string_logical().c_str(), |
| 211 | cn1.to_string().c_str(), cn1.to_string_logical().c_str())); |
| 212 | } |
| 213 | } |
| 214 | } |
| 215 | } |
| 216 | if (!valid) { |
| 217 | return CompNode::load(get_default_device()); |
| 218 | } |
| 219 | Py_XDECREF(tuple); |
| 220 | return cn; |
| 221 | } |
| 222 | |
| 223 | // Returns the dtype that would result from performing an arithmetic |
| 224 | // operation on the provided input tensors and scalars. |
no test coverage detected