mlx/docs/src/using_streams.rst
Awni Hannun cfc39d84b7
Some docs on unified memory (#62)
* doc on unified memory
2023-12-07 19:42:24 -08:00

19 lines
651 B
ReStructuredText

.. _using_streams:
Using Streams
=============
.. currentmodule:: mlx.core
Specifying the :obj:`Stream`
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
All operations (including random number generation) take an optional
keyword argument ``stream``. The ``stream`` kwarg specifies which
:obj:`Stream` the operation should run on. If the stream is unspecified then
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 :obj:`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)``.