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

Function _fastpath_getitem_cpp

imperative/python/src/tensor_utils.cpp:1145–1251  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1143}
1144
1145py::object _fastpath_getitem_cpp(py::handle inp_hdl, py::tuple tuple_val) {
1146 py::object inp = py::reinterpret_borrow<py::object>(inp_hdl);
1147 int ax = 0;
1148 bool use_ellipsis = false;
1149 size_t special_dim = 0;
1150
1151 for (size_t i = 0; i < tuple_val.size(); ++i) {
1152 PyObject* obj = tuple_val[i].ptr();
1153 if (obj == Py_Ellipsis) {
1154 use_ellipsis = true;
1155 for (size_t j = i + 1; j < tuple_val.size(); j++) {
1156 PyObject* obj_last = tuple_val[j].ptr();
1157 if (obj_last == Py_Ellipsis) {
1158 throw py::index_error("only one ellipsis is allowed.");
1159 }
1160 }
1161 }
1162 if (obj != Py_None && obj != Py_Ellipsis && obj != Py_True && obj != Py_False) {
1163 special_dim++;
1164 }
1165 }
1166
1167 size_t ndim = 0;
1168 try {
1169 ndim = getattr(inp_hdl, "ndim").cast<size_t>();
1170 } catch (py::error_already_set& err) {
1171 if (use_ellipsis) {
1172 throw py::index_error(
1173 "does not support Ellipsis when tensor's ndim is unknown.");
1174 };
1175 }
1176
1177 std::vector<std::tuple<int8_t, bool, bool, bool, bool>> cpp_items;
1178 std::vector<std::tuple<int32_t, int32_t, int32_t, int32_t>> slice_items;
1179 std::vector<int32_t> expand_items;
1180
1181 for (size_t i = 0; i < tuple_val.size(); ++i) {
1182 py::object t = tuple_val[i];
1183 if (t.ptr() == Py_Ellipsis) {
1184 ax += ndim - special_dim;
1185 } else if (PySlice_Check(t.ptr())) {
1186 PySliceObject* s = (PySliceObject*)t.ptr();
1187 std::vector<int> items;
1188 std::vector<bool> idx_items;
1189 auto push = [&](PyObject* v, int default_value) {
1190 if (v == Py_None) {
1191 items.push_back(default_value);
1192 idx_items.push_back(false);
1193 } else {
1194 auto obj = py::reinterpret_borrow<py::object>(v);
1195 items.push_back(obj.cast<int>());
1196 idx_items.push_back(true);
1197 }
1198 };
1199 push(s->start, INT_MIN);
1200 push(s->stop, INT_MAX);
1201 push(s->step, INT_MAX);
1202 if (idx_items[0] || idx_items[1] || idx_items[2]) {

Callers 1

_getitem_cppFunction · 0.85

Calls 10

pushFunction · 0.85
is_scalarFunction · 0.85
castFunction · 0.85
py_applyFunction · 0.85
resizeMethod · 0.80
makeFunction · 0.50
sizeMethod · 0.45
ptrMethod · 0.45
push_backMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected