Understanding the Relationship between Stream and Device in Computational Graphs

[khirotaka]

Hello there, great work!

Please tell me about the relationship between Stream and Device.

I recognized Stream as a so-called calculation graph where data flows, is it correct? And I thought that there was a keyword called Stream in the argument of each function to specify which device the calculation graph should be executed on, is that correct?

[awni]

That's sounds correct. This page in the docs should help. Basically every op can take a stream key word argument. You can specify a stream or a device for that argument. There are three cases:

1. If it's a stream then it runs the op on that stream.
2. If its a device then it runs the op on the default stream for the given device.
3. If its unspecified then it runs the op on the default stream of the default devices.

In terms of scheduling and dependencies, MLX manages dependencies between streams. So if you have an array output from one stream that is input to another, the second stream will wait until the array is ready from the first.