MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / _make_shape_tuple

Function _make_shape_tuple

imperative/python/src/tensor_utils.cpp:331–375  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

329}
330
331py::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
377bool is_tensor(py::handle arg) {
378 return bool(TensorWrapper::try_cast(arg.ptr()));

Callers 5

_try_cond_takeFunction · 0.85
_expand_bool_dimFunction · 0.85
_reshape_cppFunction · 0.85
_adaptive_pool2d_cppFunction · 0.85
tensor_utils.cppFile · 0.85

Calls 4

try_castFunction · 0.85
ptrMethod · 0.45
sizeMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected