zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Adds device context manager (#679)

Browse Source
This commit is contained in:
Diogo
2024-02-14 17:14:58 -05:00
committed by GitHub
parent ccf1645995
commit 35431a4ac8
15 changed files with 230 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
mlx/utils.cpp
View File

@@ -7,6 +7,16 @@
namespace mlx::core {
Stream to_stream(StreamOrDevice s) {
if (std::holds_alternative<std::monostate>(s)) {
return default_stream(default_device());
} else if (std::holds_alternative<Device>(s)) {
return default_stream(std::get<Device>(s));
} else {
return std::get<Stream>(s);
}
}
void PrintFormatter::print(std::ostream& os, bool val) {
if (capitalize_bool) {
os << (val ? "True" : "False");