diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp
index 32bdf0c0..75327fdb 100644
--- a/mlx/scheduler.cpp
+++ b/mlx/scheduler.cpp
@@ -16,6 +16,10 @@ void synchronize(Stream s) {
   }
 }
 
+void synchronize(ThreadLocalStream s) {
+  synchronize(stream_from_thread_local_stream(s));
+}
+
 void synchronize() {
   synchronize(default_stream(default_device()));
 }
diff --git a/mlx/stream.h b/mlx/stream.h
index ae333e03..fd938955 100644
--- a/mlx/stream.h
+++ b/mlx/stream.h
@@ -49,6 +49,9 @@ MLX_API void synchronize();
 /* Synchronize with the provided stream. */
 MLX_API void synchronize(Stream);
 
+/* Synchronize with the stream corresponding to the current thread. */
+MLX_API void synchronize(ThreadLocalStream);
+
 /* Destroy all streams created in current thread. */
 MLX_API void clear_streams();
 
diff --git a/python/src/stream.cpp b/python/src/stream.cpp
index 8764eaf7..62029839 100644
--- a/python/src/stream.cpp
+++ b/python/src/stream.cpp
@@ -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");