| 454 | } |
| 455 | |
| 456 | mx::array mlx_get_item(const mx::array& src, const nb::object& obj) { |
| 457 | if (nb::isinstance<nb::slice>(obj)) { |
| 458 | return mlx_get_item_slice(src, nb::cast<nb::slice>(obj)); |
| 459 | } else if (nb::isinstance<mx::array>(obj)) { |
| 460 | return mlx_get_item_array(src, nb::cast<mx::array>(obj)); |
| 461 | } else if (is_index_scalar(obj)) { |
| 462 | return mlx_get_item_int(src, obj); |
| 463 | } else if (nb::isinstance<nb::tuple>(obj)) { |
| 464 | return mlx_get_item_nd(src, nb::cast<nb::tuple>(obj)); |
| 465 | } else if (nb::isinstance<nb::ellipsis>(obj)) { |
| 466 | return src; |
| 467 | } else if (obj.is_none()) { |
| 468 | return expand_dims(src, 0); |
| 469 | } else if (nb::isinstance<nb::list>(obj)) { |
| 470 | return mlx_get_item_array( |
| 471 | src, array_from_list(nb::cast<nb::list>(obj), {})); |
| 472 | } |
| 473 | throw std::invalid_argument("Cannot index mlx array using the given type."); |
| 474 | } |
| 475 | |
| 476 | std::tuple<std::vector<mx::array>, mx::array, std::vector<int>> |
| 477 | mlx_scatter_args_int( |
nothing calls this directly
no test coverage detected