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

Function _remove_ellipsis

imperative/python/src/tensor_utils.cpp:637–700  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

635}
636
637py::tuple _remove_ellipsis(py::object tensor, py::tuple tuple_val) {
638 size_t tuple_size = tuple_val.size();
639 size_t ndim_sum = 0, cur_sum = 0;
640 int pos = -1;
641 bool has_unknown_ndim_bool_index = false;
642 for (size_t i = 0; i < tuple_size; ++i) {
643 py::object handle = tuple_val[i];
644 if (handle.is_none()) {
645 continue;
646 } else if (handle.ptr() == Py_Ellipsis) {
647 pos = static_cast<int>(i);
648 for (size_t j = 0; j < i; ++j) {
649 py::object t = tuple_val[j];
650 if (t.ptr() == Py_Ellipsis) {
651 throw py::index_error("only one ellipsis is allowed.");
652 }
653 }
654 } else {
655 size_t ndim_incr = 1;
656 if (hasattr(handle, "dtype") && is_bool_dtype(handle.ptr()) &&
657 hasattr(handle, "ndim")) {
658 py::object ndim;
659 try {
660 ndim = getattr(handle, "ndim");
661 } catch (py::error_already_set& err) {
662 has_unknown_ndim_bool_index = true;
663 }
664 if (PyLong_Check(ndim.ptr())) {
665 ndim_incr = PyLong_AsLong(ndim.ptr());
666 } else {
667 has_unknown_ndim_bool_index = true;
668 }
669 }
670 cur_sum += ndim_incr;
671 }
672 }
673 if (pos == -1) {
674 return tuple_val;
675 } else {
676 if (has_unknown_ndim_bool_index) {
677 throw py::index_error(
678 "does not support bool index with unknown shape when using "
679 "Ellipsis.");
680 }
681 try {
682 ndim_sum = getattr(tensor, "ndim").cast<size_t>();
683 } catch (py::error_already_set& err) {
684 throw py::index_error(
685 "does not support Ellipsis when tensor's ndim is unknown.");
686 }
687 py::tuple ret(ndim_sum - cur_sum + tuple_size - 1);
688 size_t idx = 0;
689 for (size_t i = 0; i < tuple_size; ++i) {
690 if (i == pos) {
691 for (size_t j = cur_sum; j < ndim_sum; ++j) {
692 ret[idx++] = PySlice_New(NULL, NULL, NULL);
693 }
694 } else {

Callers 1

_unpack_indexesFunction · 0.85

Calls 3

is_bool_dtypeFunction · 0.85
sizeMethod · 0.45
ptrMethod · 0.45

Tested by

no test coverage detected