MLX Data

MLX Data is a framework-agnostic data loading library developed by Apple Machine Learning Research. It is designed to facilitate high-performance data loading for machine learning training or standalone data pre-processing.

Key characteristics:

• Framework Agnostic: Compatible with PyTorch, JAX, or MLX.
• Multi-threaded: Leverages multiple threads for data processing pipelines, avoiding the complexity of multi-process management or symbolic languages.
• Python-centric: You can use standard Python to implement processing logic, handle data transformations, or cause side effects directly within the pipeline.

The benchmarks/comparative directory contains performance comparisons between mlx.data, PyTorch DataLoaders, and tf.data. These benchmarks aim to demonstrate the conciseness and speed of mlx.data.

Note that these are comparative benchmarks and may not reflect real-world performance for your specific use case.

A Stream acts like an iterator. When you compose streams (e.g., by creating evens from numbers), the new stream is a pointer to the underlying source. Advancing the base stream will also advance any derived streams.

If the underlying data source supports it, you can use .reset() to restart the stream from the beginning.

import mlx.data as dx

numbers = dx.stream_python_iterable(lambda: ({"x": i} for i in range(10)))
evens = numbers.sample_transform(lambda s: s if s["x"] % 2 == 0 else dict())

# Advancing 'numbers' affects 'evens'
print(next(numbers)) # {'x': 0}
print(next(evens))   # {'x': 2} (skips 1 because it was filtered)

# Resetting the stream
evens.reset()
print(next(evens))   # {'x': 0}

A Buffer in mlx.data is an indexable container of samples. It behaves similarly to a Python list, allowing for random access and transformations that require knowing the total size or order of the data. You can use key_transform to apply a function to a specific key across all samples in the buffer.

import mlx.data as dx

# Create a buffer from a list of dictionaries
numbers = dx.buffer_from_vector([{"x": i} for i in range(10)])

# Apply a transformation to the key 'x'
evens = numbers.key_transform("x", lambda x: 2*x)

print(evens)
# prints Buffer(size=10, keys={'x'})

print(evens[3])
# prints {'x': array(6)}

print(len(evens))
# prints 10

In MLX Data, a sample is the fundamental unit of data. It is a dictionary that maps string keys to array values.

• In Python: A dictionary where keys are strings and values are anything implementing the buffer protocol (e.g., NumPy arrays, scalars, or bytes).
• In C++: An instance of mlx::data::Sample, which is a std::unordered_map<std::string, std::shared_ptr<mlx::data::Array>>.

Commonly used types for values include:

• Arrays: e.g., np.array(0).
• Scalars: These are automatically cast to scalar arrays.
• Strings: Represented as Unicode.
• Bytes: Recommended for paths or raw data (e.g., b"path/to/file").

# Valid samples
sample = {"hello": np.array(0)}
sample = {"scalar": 42}
sample = {"key": "value"}

# Recommended way to handle paths/bytes
sample = {"key": b"path/to/my/file"}
sample = {"key": "value".encode("ascii")}

To avoid complex redirection logic when configuring pipelines based on command-line arguments or configuration settings, mlx.data provides conditional variants for almost all transformation methods.

Every transformation method has a corresponding *_if(cond: bool, *args, **kwargs) variant. This allows you to write linear, top-to-bottom pipelines where an operation is only applied if the provided cond evaluates to True.

# Assuming we have a buffer with image files and labels in dset
dset = (
    dset
    .load_image("image_file", output_key="image")
    .image_random_crop_if(enable_random_crop, "image", 256, 256)
    .image_random_h_flip_if(flip_prob > 0, "image", flip_prob)
    .key_transform_if(brightness_range > 0, "image",
                      lambda x: ((1 + brightness_range * np.random.rand(x.shape[:2])[..., None]) * x).astype(x.dtype))
)

Choosing between a Buffer and a Stream depends on your data access requirements: