mlx.data.Stream.sliding_window

mlx.data.Stream.sliding_window#

Stream.sliding_window(self: mlx.data._c.Stream, key: str, size: int, stride: int, dim: int = -1, index_key: str = '') → mlx.data._c.Stream#

Creates sample by sliding a window over the array at key.

Commonly used in sequence processing pipelines to deal with very larger documents.

import mlx.data as dx

dset = dx.buffer_from_vector({"x": np.arange(10), "unchanged_keys": 10}).to_stream()

for sample in dset.sliding_window("x", 3, 2):
    print(sample["x"])

# prints
# [0, 1, 2]
# [2, 3, 4]
# [4, 5, 6]
# [6, 7, 8]
# [8, 9]
Parameters:

  • key (str) – The sample key that contains the array we are operating on.

  • size (int) – The size of the sliding window.

  • stride (int) – The stride of the sliding window.

  • dim (int) – Which dimension are we sliding the window over. (default: -1)

  • index_key (str) – If provided, store the index of the sliding window in that key. (default: “”)

Contents