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

Function _expand_bool_dim

imperative/python/src/tensor_utils.cpp:704–776  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

702py::object _reshape_cpp(py::handle inp_hdl, py::handle args);
703
704py::tuple _expand_bool_dim(py::object tensor, py::tuple tuple_val) {
705 py::tuple cur_shape = _make_shape_tuple(py::handle(getattr(tensor, "shape")));
706 py::list new_tuple_val(0);
707
708 size_t offset = 0;
709 size_t tdim = 0;
710 size_t nonedim = 0;
711 for (size_t i = 0; i < tuple_val.size(); ++i) {
712 py::handle k = tuple_val[i];
713 if (k.ptr() == Py_None) {
714 nonedim++;
715 new_tuple_val.append(k);
716 continue;
717 }
718 if (is_bool_dtype(k.ptr())) {
719 size_t ndim = getattr(k, "ndim").cast<size_t>();
720 if (ndim > 1) {
721 py::tuple ishape = _make_shape_tuple(py::handle(getattr(k, "shape")));
722 for (size_t j = 0; j < ndim; ++j) {
723 if (cur_shape[tdim + j - offset].cast<size_t>() !=
724 ishape[j].cast<size_t>()) {
725 std::string msg =
726 "boolean index did not match tensor along "
727 "dimension " +
728 std::to_string(tdim + j) + "; dimension is " +
729 std::to_string(
730 cur_shape[tdim + j - offset].cast<size_t>()) +
731 " but corresponding boolean dimension is " +
732 std::to_string(ishape[j].cast<size_t>());
733 throw py::index_error(msg.c_str());
734 }
735 }
736 py::object new_k = getattr(k, "reshape")(-1);
737 py::object kshape = getattr(new_k, "shape");
738 py::list new_shape(0);
739 PyObject* sym = PyObject_CallObject(cpp_use_symbolic_shape, nullptr);
740 bool is_sym = (sym == Py_True);
741 Py_XDECREF(sym);
742 if (is_sym) {
743 py::object tshape = getattr(tensor, "shape");
744 for (size_t j = 0; j < i - nonedim; ++j) {
745 new_shape.append(tshape[py::int_(j)]);
746 }
747 new_shape.append(kshape[py::int_(0)]);
748 for (size_t j = tdim + ndim - offset; j < cur_shape.size(); ++j) {
749 new_shape.append(cur_shape[j]);
750 }
751 py::object shape_tensor = _astensor1d_cpp(
752 new_shape, py::none(), py::none(), py::none());
753 tensor = _reshape_cpp(tensor, shape_tensor);
754 cur_shape = _make_shape_tuple(shape_tensor);
755 } else {
756 for (size_t j = 0; j < i - nonedim; ++j) {
757 new_shape.append(cur_shape[j]);
758 }
759 new_shape.append(py::reinterpret_borrow<py::tuple>(kshape)[0]);
760 for (size_t j = tdim + ndim - offset; j < cur_shape.size(); ++j) {
761 new_shape.append(cur_shape[j]);

Callers 1

_unpack_indexesFunction · 0.85

Calls 8

_make_shape_tupleFunction · 0.85
is_bool_dtypeFunction · 0.85
to_stringFunction · 0.85
_astensor1d_cppFunction · 0.85
_reshape_cppFunction · 0.85
sizeMethod · 0.45
ptrMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected