Adds device context manager (#679)

python/src/stream.cpp

@@ -12,7 +12,12 @@ using namespace py::literals;
 using namespace mlx::core;
 
 void init_stream(py::module_& m) {
-  py::class_<Stream>(m, "Stream")
+  py::class_<Stream>(
+      m,
+      "Stream",
+      R"pbdoc(
+      A stream for running operations on a given device.
+      )pbdoc")
       .def(py::init<int, Device>(), "index"_a, "device"_a)
       .def_readonly("device", &Stream::device)
       .def(
@@ -28,7 +33,27 @@ void init_stream(py::module_& m) {
 
   py::implicitly_convertible<Device::DeviceType, Device>();
 
-  m.def("default_stream", &default_stream, "device"_a);
-  m.def("set_default_stream", &set_default_stream, "stream"_a);
-  m.def("new_stream", &new_stream, "device"_a);
+  m.def(
+      "default_stream",
+      &default_stream,
+      "device"_a,
+      R"pbdoc(Get the device's default stream.)pbdoc");
+  m.def(
+      "set_default_stream",
+      &set_default_stream,
+      "stream"_a,
+      R"pbdoc(
+        Set the default stream.
+
+        This will make the given stream the default for the
+        streams device. It will not change the default device.
+
+        Args:
+          stream (stream): Stream to make the default.
+      )pbdoc");
+  m.def(
+      "new_stream",
+      &new_stream,
+      "device"_a,
+      R"pbdoc(Make a new stream on the given device.)pbdoc");
 }