Fix synchronize for ThreadLocalStream (#3429)

2026-04-20 11:29:49 -07:00
parent b4ddf9b374
commit 705c828feb
3 changed files with 15 additions and 3 deletions
@@ -162,15 +162,20 @@ void init_stream(nb::module_& m) {
      )pbdoc");
  m.def(
      "synchronize",
-      [](const std::optional<mx::Stream>& s) {
-        s ? mx::synchronize(s.value()) : mx::synchronize();
+      [](mx::StreamOrDevice s) {
+        if (std::holds_alternative<std::monostate>(s)) {
+          mx::synchronize();
+        } else {
+          mx::synchronize(mx::to_stream(s));
+        }
      },
      "stream"_a = nb::none(),
      R"pbdoc(
      Synchronize with the given stream.

      Args:
-        stream (Stream, optional): The stream to synchronize with. If ``None``
+        stream (Stream, optional): Stream to synchronize. If device is
+           provided the default stream for that device is used. If ``None``
           then the default stream of the default device is used.
           Default: ``None``.
      )pbdoc");