| 612 | } |
| 613 | |
| 614 | py::tuple _try_cond_take(py::handle tensor, py::handle index) { |
| 615 | if (!hasattr(index, "dtype") || !hasattr(index, "shape")) { |
| 616 | return py::tuple(); |
| 617 | } |
| 618 | if (!is_bool_dtype(index.ptr()) || |
| 619 | _make_shape_tuple(getattr(index, "shape")) |
| 620 | .not_equal(_make_shape_tuple(getattr(tensor, "shape")))) { |
| 621 | return py::tuple(); |
| 622 | } |
| 623 | py::object iobj; |
| 624 | if (PyArray_Check(index.ptr())) { |
| 625 | iobj = _Const( |
| 626 | index, py::cast((mgb::DType)dtype::Bool()), getattr(tensor, "device")); |
| 627 | } else { |
| 628 | iobj = py::reinterpret_borrow<py::object>(index); |
| 629 | } |
| 630 | std::shared_ptr<OpDef> op = CondTake::make(); |
| 631 | py::object Op = py::cast(op); |
| 632 | PyObject* p[3] = {Op.ptr(), tensor.ptr(), iobj.ptr()}; |
| 633 | py::tuple ret = py::reinterpret_steal<py::object>(py_apply(NULL, p, 3)); |
| 634 | return ret; |
| 635 | } |
| 636 | |
| 637 | py::tuple _remove_ellipsis(py::object tensor, py::tuple tuple_val) { |
| 638 | size_t tuple_size = tuple_val.size(); |
no test coverage detected