Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/ml-explore/mlx/llms.txt

Use this file to discover all available pages before exploring further.

All operations in MLX (including random number generation) take an optional keyword argument stream. The stream kwarg specifies which stream the operation should run on.

​
Specifying the Stream

If the stream is unspecified, the operation is run on the default stream of the default device: 
mx.default_stream(mx.default_device())
The stream kwarg can also be a Device (e.g., stream=my_device), in which case the operation is run on the default stream of the provided device: 
mx.default_stream(my_device)

​
Example Usage

import mlx.core as mx

# Create a custom stream
my_stream = mx.Stream(mx.cpu)

# Run an operation on the custom stream
result = mx.add(a, b, stream=my_stream)

# Or use a device directly
result = mx.add(a, b, stream=mx.gpu)
Streams allow for asynchronous execution and can help improve performance by overlapping computation and data transfer.
When working with multiple devices or optimizing performance, explicitly specifying streams can give you fine-grained control over operation scheduling.