| 587 | } |
| 588 | |
| 589 | py::object _get_index(py::object tensor, py::object src) { |
| 590 | if (!TensorWrapper::try_cast(tensor.ptr())) { |
| 591 | auto get_const = [&](mgb::DType dtype) -> py::object { |
| 592 | return _Const(tensor, py::cast(dtype), src.attr("device")); |
| 593 | }; |
| 594 | if (is_bool_list(tensor.ptr()) || is_bool_dtype(tensor.ptr())) { |
| 595 | tensor = get_const(dtype::Bool()); |
| 596 | } else { |
| 597 | tensor = get_const(dtype::Int32()); |
| 598 | } |
| 599 | if (!is_bool_dtype(tensor.ptr())) { |
| 600 | return tensor; |
| 601 | } |
| 602 | } else { |
| 603 | if (!is_bool_dtype(tensor.ptr())) { |
| 604 | return tensor; |
| 605 | } |
| 606 | } |
| 607 | std::shared_ptr<OpDef> op = CondTake::make(); |
| 608 | py::object Op = py::cast(op); |
| 609 | PyObject* p[3] = {Op.ptr(), tensor.ptr(), tensor.ptr()}; |
| 610 | py::tuple ret = py::reinterpret_steal<py::object>(py_apply(NULL, p, 3)); |
| 611 | return ret[1]; |
| 612 | } |
| 613 | |
| 614 | py::tuple _try_cond_take(py::handle tensor, py::handle index) { |
| 615 | if (!hasattr(index, "dtype") || !hasattr(index, "shape")) { |
no test coverage detected