Add clear_streams API for cleanup before exit (#3395)

2026-04-14 18:41:32 +09:00
parent 520cea2bec
commit 435f0b6cdb
10 changed files with 53 additions and 3 deletions
@@ -133,6 +133,10 @@ void init_stream(nb::module_& m) {
      &mx::new_stream,
      "device"_a,
      R"pbdoc(Make a new stream on the given device.)pbdoc");
+  m.def(
+      "clear_streams",
+      &mx::clear_streams,
+      R"pbdoc(Destroy all streams created in current thread.)pbdoc");

  nb::class_<PyStreamContext>(m, "StreamContext", R"pbdoc(
        A context manager for setting the current device and stream.