| 1441 | } |
| 1442 | |
| 1443 | py::object _split_cpp( |
| 1444 | py::handle inp_hdl, py::handle nsplits_or_sections_hdl, py::handle axis_hdl) { |
| 1445 | py::object shape_obj = getattr(inp_hdl, "shape"); |
| 1446 | py::object n_total = shape_obj[axis_hdl]; |
| 1447 | int ndim = shape_obj.attr("__len__")().cast<int>(); |
| 1448 | int axis = axis_hdl.cast<int>(); |
| 1449 | if (axis >= ndim) { |
| 1450 | throw py::value_error("Invalid axis " + std::to_string(axis)); |
| 1451 | } |
| 1452 | int n_sections; |
| 1453 | bool is_array; |
| 1454 | if (is_py_sequence(nsplits_or_sections_hdl)) { |
| 1455 | n_sections = PySequence_Length(nsplits_or_sections_hdl.ptr()) + 1; |
| 1456 | is_array = true; |
| 1457 | } else { |
| 1458 | n_sections = getattr(nsplits_or_sections_hdl, "__int__")().cast<int>(); |
| 1459 | is_array = false; |
| 1460 | } |
| 1461 | py::list partitions; |
| 1462 | std::shared_ptr<OpDef> op; |
| 1463 | std::vector<PyObject*> p; |
| 1464 | if (is_array) { |
| 1465 | py::list div_points; |
| 1466 | py::list sections = py::reinterpret_borrow<py::object>(nsplits_or_sections_hdl); |
| 1467 | div_points.append(0); |
| 1468 | for (size_t i = 0; i < sections.size(); ++i) { |
| 1469 | div_points.append(sections[i]); |
| 1470 | } |
| 1471 | div_points.append(n_total); |
| 1472 | for (size_t i = 1; i < div_points.size(); ++i) { |
| 1473 | if (div_points[i - 1] > div_points[i]) { |
| 1474 | throw py::value_error( |
| 1475 | "Invalid nsplits_or_secions: " + |
| 1476 | repr(nsplits_or_sections_hdl).cast<std::string>()); |
| 1477 | } |
| 1478 | py::object pos = div_points[i] - div_points[i - 1]; |
| 1479 | if (is_tensor(pos)) { |
| 1480 | partitions.append(pos); |
| 1481 | } else { |
| 1482 | partitions.append( |
| 1483 | _Const(pos, py::cast((mgb::DType)dtype::Int32()), |
| 1484 | getattr(inp_hdl, "device"))); |
| 1485 | } |
| 1486 | } |
| 1487 | op = Split::make(axis, 0); |
| 1488 | p.resize(partitions.size() + 2); |
| 1489 | for (size_t i = 0; i < partitions.size(); ++i) { |
| 1490 | p[i + 2] = partitions[i].ptr(); |
| 1491 | } |
| 1492 | } else { |
| 1493 | if (n_sections <= 0) { |
| 1494 | throw py::value_error("Number sections must be larger than 0"); |
| 1495 | } |
| 1496 | if (py::int_(n_sections) > n_total) { |
| 1497 | throw py::value_error( |
| 1498 | "The size " + repr(n_total).cast<std::string>() + " at dim " + |
| 1499 | std::to_string(axis) + " cannot be split into " + |
| 1500 | std::to_string(n_sections) + " sections"); |
no test coverage detected