Add synchronize function (#1006)

* add synchronize function

* fix linux

* fix linux

* fix and fix docs

* fix test

* try synchronize in stream destroy

* synchronize works for both cpu and gpu

python/src/metal.cpp

@@ -2,6 +2,7 @@
 
 #include "mlx/backend/metal/metal.h"
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 
 namespace nb = nanobind;
@@ -99,9 +100,6 @@ void init_metal(nb::module_& m) {
       Args:
         path (str): The path to save the capture which should have
           the extension ``.gputrace``.
-
-      Returns:
-        bool: Whether the capture was successfully started.
       )pbdoc");
   metal.def(
       "stop_capture",

python/src/stream.cpp

@@ -129,4 +129,17 @@ void init_stream(nb::module_& m) {
               # Operations here will use mx.cpu by default.
               pass
       )pbdoc");
+  m.def(
+      "synchronize",
+      [](const std::optional<Stream>& s) {
+        s ? synchronize(s.value()) : synchronize();
+      },
+      "stream"_a = nb::none(),
+      R"pbdoc(
+      Synchronize with the given stream.
+
+      Args:
+        (Stream, optional): The stream to synchronize with. If ``None`` then
+           the default stream of the default device is used. Default: ``None``.
+      )pbdoc");
 }