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

Function get_slice_params

python/src/indexing.cpp:55–72  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

53}
54
55void get_slice_params(
56 mx::ShapeElem& starts,
57 mx::ShapeElem& ends,
58 mx::ShapeElem& strides,
59 const nb::slice& in_slice,
60 int axis_size) {
61 // Following numpy's convention
62 // Assume n is the number of elements in the dimension being sliced.
63 // Then, if i is not given it defaults to 0 for k > 0 and n - 1 for
64 // k < 0 . If j is not given it defaults to n for k > 0 and -n-1 for
65 // k < 0 . If k is not given it defaults to 1
66
67 strides = get_slice_int(nb::getattr(in_slice, "step"), 1);
68 starts = get_slice_int(
69 nb::getattr(in_slice, "start"), strides < 0 ? axis_size - 1 : 0);
70 ends = get_slice_int(
71 nb::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
72}
73
74mx::array get_int_index(nb::object idx, int axis_size) {
75 int idx_ = safe_to_int32(idx);

Callers 6

mlx_get_item_sliceFunction · 0.85
mlx_gather_ndFunction · 0.85
mlx_get_item_ndFunction · 0.85
mlx_scatter_args_sliceFunction · 0.85
mlx_scatter_args_ndFunction · 0.85

Calls 1

get_slice_intFunction · 0.85

Tested by

no test coverage detected