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

Function _setitem_cpp

imperative/python/src/tensor_utils.cpp:1321–1441  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1319}
1320
1321py::object _setitem_cpp(py::handle inp_hdl, py::handle idx_hdl, py::handle val_hdl) {
1322 py::object org_shape = getattr(inp_hdl, "shape");
1323 py::object val = py::reinterpret_borrow<py::object>(val_hdl);
1324 bool is_val_scalar = false;
1325 float value;
1326 if (PyLong_Check(val.ptr())) {
1327 is_val_scalar = true;
1328 value = static_cast<float>(PyLong_AsDouble(val.ptr()));
1329 }
1330 if (PyFloat_Check(val.ptr())) {
1331 is_val_scalar = true;
1332 value = static_cast<float>(PyFloat_AsDouble(val.ptr()));
1333 }
1334 if (TensorWrapper::try_cast(idx_hdl.ptr()) && is_bool_dtype(idx_hdl.ptr()) &&
1335 is_val_scalar && enable_fastpath(inp_hdl)) {
1336 std::vector<PyObject*> q(3);
1337 std::shared_ptr<OpDef> Op = MaskedFill::make(value);
1338 py::object maskedfill = py::cast(Op);
1339 q[0] = maskedfill.ptr();
1340 q[1] = inp_hdl.ptr();
1341 q[2] = idx_hdl.ptr();
1342 py::tuple result =
1343 py::reinterpret_steal<py::object>(py_apply(NULL, q.data(), q.size()));
1344 py::object res = result[0];
1345 return res;
1346 }
1347 if (!TensorWrapper::try_cast(val.ptr())) {
1348 val = _Const(val_hdl, getattr(inp_hdl, "dtype"), getattr(inp_hdl, "device"));
1349 }
1350
1351 py::tuple tuple_val;
1352 if (py::isinstance<py::tuple>(idx_hdl)) {
1353 tuple_val = py::reinterpret_borrow<py::tuple>(idx_hdl);
1354 } else {
1355 tuple_val = py::make_tuple(idx_hdl);
1356 }
1357 py::tuple up = _unpack_indexes(inp_hdl, tuple_val);
1358 py::object tensor = py::reinterpret_borrow<py::object>(up[0]);
1359 py::list tensors = py::reinterpret_borrow<py::list>(up[1]);
1360 py::list py_items = py::reinterpret_borrow<py::list>(up[2]);
1361 std::vector<std::tuple<int8_t, bool, bool, bool, bool>> cpp_items;
1362 std::vector<std::tuple<int32_t, int32_t, int32_t, int32_t>> slice_items;
1363 for (size_t i = 0; i < py_items.size(); ++i) {
1364 py::list item = py::reinterpret_borrow<py::list>(py_items[i]);
1365 cpp_items.push_back(
1366 {item[0].cast<int8_t>(), item[1].cast<bool>(), item[2].cast<bool>(),
1367 item[3].cast<bool>(), item[4].cast<bool>()});
1368 }
1369 std::shared_ptr<OpDef> op, set_op;
1370 if (up[3].cast<bool>()) {
1371 op = Subtensor::make(cpp_items, slice_items);
1372 } else {
1373 op = IndexingMultiAxisVec::make(cpp_items);
1374 }
1375 std::vector<PyObject*> p;
1376 p.resize(tensors.size() + 2);
1377 py::object Op = py::cast(op);
1378 p[0] = Op.ptr();

Callers 1

tensor_utils.cppFile · 0.85

Calls 15

try_castFunction · 0.85
is_bool_dtypeFunction · 0.85
enable_fastpathFunction · 0.85
castFunction · 0.85
py_applyFunction · 0.85
_ConstFunction · 0.85
_unpack_indexesFunction · 0.85
to_stringFunction · 0.85
_broadcast_cppFunction · 0.85
_reshape_cppFunction · 0.85
resizeMethod · 0.80
attrMethod · 0.80

Tested by

no test coverage detected