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())
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.