| 329 | } |
| 330 | |
| 331 | py::tuple _make_shape_tuple(py::handle shape) { |
| 332 | py::list orig; |
| 333 | py::list ret(0); |
| 334 | auto solve_one = [&](py::handle val) { |
| 335 | if (TensorWrapper::try_cast(val.ptr())) { |
| 336 | py::object np = getattr(val, "numpy")(); |
| 337 | PyArrayObject* arr = (PyArrayObject*)np.ptr(); |
| 338 | PyObject* maybe_list = PyArray_ToList(arr); |
| 339 | if (PyList_Check(maybe_list)) { |
| 340 | py::list may = py::reinterpret_steal<py::list>(maybe_list); |
| 341 | for (size_t i = 0; i < may.size(); ++i) { |
| 342 | ret.append(may[i]); |
| 343 | } |
| 344 | } else { |
| 345 | mgb_assert(PyLong_Check(maybe_list)); |
| 346 | ret.append(PyLong_AsLong(maybe_list)); |
| 347 | Py_XDECREF(maybe_list); |
| 348 | } |
| 349 | } else if (PyArray_Check(val.ptr())) { |
| 350 | ret.append(PyArray_PyIntAsInt(val.ptr())); |
| 351 | } else { |
| 352 | ret.append(PyLong_AsLong(val.ptr())); |
| 353 | } |
| 354 | }; |
| 355 | if (PyArray_Check(shape.ptr()) && !PyArray_CheckAnyScalar(shape.ptr())) { |
| 356 | orig = py::reinterpret_steal<py::list>( |
| 357 | PyArray_ToList((PyArrayObject*)shape.ptr())); |
| 358 | for (size_t i = 0; i < orig.size(); ++i) { |
| 359 | solve_one(orig[i]); |
| 360 | } |
| 361 | } else if (PyList_Check(shape.ptr())) { |
| 362 | orig = py::reinterpret_borrow<py::list>(shape); |
| 363 | for (size_t i = 0; i < orig.size(); ++i) { |
| 364 | solve_one(orig[i]); |
| 365 | } |
| 366 | } else if (PyTuple_Check(shape.ptr())) { |
| 367 | py::tuple tup = py::reinterpret_borrow<py::tuple>(shape); |
| 368 | for (size_t i = 0; i < tup.size(); ++i) { |
| 369 | solve_one(tup[i]); |
| 370 | } |
| 371 | } else { |
| 372 | solve_one(shape); |
| 373 | } |
| 374 | return py::reinterpret_steal<py::tuple>(PyList_AsTuple(ret.ptr())); |
| 375 | } |
| 376 | |
| 377 | bool is_tensor(py::handle arg) { |
| 378 | return bool(TensorWrapper::try_cast(arg.ptr())); |
no test coverage detected