MCPcopy Create free account
hub / github.com/ml-explore/mlx / mlx_get_item

Function mlx_get_item

python/src/indexing.cpp:456–474  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

454}
455
456mx::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
476std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
477mlx_scatter_args_int(

Callers

nothing calls this directly

Calls 7

mlx_get_item_sliceFunction · 0.85
mlx_get_item_arrayFunction · 0.85
is_index_scalarFunction · 0.85
mlx_get_item_intFunction · 0.85
mlx_get_item_ndFunction · 0.85
expand_dimsFunction · 0.85
array_from_listFunction · 0.85

Tested by

no test coverage detected