mlx.core.slice

目录

mlx.core.slice#

slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: None | Stream | Device = None) array#

Extract a sub-array from the input array.

参数:
  • a (array) -- Input array

  • start_indices (array) -- The index location to start the slice at.

  • axes (tuple(int)) -- The axes corresponding to the indices in start_indices.

  • slice_size (tuple(int)) -- The size of the slice.

返回:

The sliced output array.

返回类型:

array

示例

>>> a = mx.array([[1, 2, 3], [4, 5, 6]])
>>> mx.slice(a, start_indices=mx.array(1), axes=(0,), slice_size=(1, 2))
array([[4, 5]], dtype=int32)
>>>
>>> mx.slice(a, start_indices=mx.array(1), axes=(1,), slice_size=(2, 1))
array([[2],
       [5]], dtype=int32)