| 42 | }; |
| 43 | |
| 44 | void init_stream(nb::module_& m) { |
| 45 | nb::class_<mx::Stream>( |
| 46 | m, |
| 47 | "Stream", |
| 48 | R"pbdoc( |
| 49 | A stream for running operations on a given device. |
| 50 | )pbdoc") |
| 51 | .def_ro("device", &mx::Stream::device) |
| 52 | .def( |
| 53 | "__repr__", |
| 54 | [](const mx::Stream& s) { |
| 55 | std::ostringstream os; |
| 56 | os << s; |
| 57 | return os.str(); |
| 58 | }) |
| 59 | .def("__eq__", [](const mx::Stream& s, const nb::object& other) { |
| 60 | return nb::isinstance<mx::Stream>(other) && |
| 61 | s == nb::cast<mx::Stream>(other); |
| 62 | }); |
| 63 | |
| 64 | nb::class_<mx::ThreadLocalStream>( |
| 65 | m, |
| 66 | "ThreadLocalStream", |
| 67 | R"pbdoc( |
| 68 | A stream that will be unique per thread and can be used to run operations on a given device. |
| 69 | )pbdoc") |
| 70 | .def_ro("device", &mx::ThreadLocalStream::device) |
| 71 | .def( |
| 72 | "__repr__", |
| 73 | [](const mx::ThreadLocalStream& s) { |
| 74 | std::ostringstream os; |
| 75 | os << "ThreadLocalStream(" << s.device << ", " << s.index << ")"; |
| 76 | return os.str(); |
| 77 | }) |
| 78 | .def( |
| 79 | "__eq__", |
| 80 | [](const mx::ThreadLocalStream& s, const nb::object& other) { |
| 81 | return nb::isinstance<mx::ThreadLocalStream>(other) && |
| 82 | s == nb::cast<mx::ThreadLocalStream>(other); |
| 83 | }); |
| 84 | |
| 85 | nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>(); |
| 86 | |
| 87 | m.def( |
| 88 | "default_stream", |
| 89 | &mx::default_stream, |
| 90 | "device"_a, |
| 91 | R"pbdoc(Get the device's default stream.)pbdoc"); |
| 92 | m.def( |
| 93 | "set_default_stream", |
| 94 | &mx::set_default_stream, |
| 95 | "stream"_a, |
| 96 | R"pbdoc( |
| 97 | Set the default stream. |
| 98 | |
| 99 | This will make the given stream the default for the |
| 100 | streams device. It will not change the default device. |
| 101 |
no test coverage detected