| 926 | } |
| 927 | |
| 928 | std::optional<mx::array> extract_boolean_mask(const nb::object& obj) { |
| 929 | using NDArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>; |
| 930 | if (nb::isinstance<nb::bool_>(obj)) { |
| 931 | return mx::array(nb::cast<bool>(obj), mx::bool_); |
| 932 | } else if (nb::isinstance<mx::array>(obj)) { |
| 933 | auto mask = nb::cast<mx::array>(obj); |
| 934 | if (mask.dtype() == mx::bool_) { |
| 935 | return mask; |
| 936 | } |
| 937 | } else if (nb::isinstance<NDArray>(obj)) { |
| 938 | auto mask = nb::cast<NDArray>(obj); |
| 939 | if (mask.dtype() == nb::dtype<bool>()) { |
| 940 | return nd_array_to_mlx(mask, mx::bool_); |
| 941 | } |
| 942 | } else if (nb::isinstance<nb::list>(obj)) { |
| 943 | auto mask = array_from_list(nb::cast<nb::list>(obj), {}); |
| 944 | if (mask.dtype() == mx::bool_) { |
| 945 | return mask; |
| 946 | } |
| 947 | } |
| 948 | return std::nullopt; |
| 949 | } |
| 950 | |
| 951 | void mlx_set_item( |
| 952 | mx::array& src, |
no test coverage detected