forked from ml-explore/mlx
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathprint.cpp
More file actions
86 lines (71 loc) · 2.56 KB
/
Copy pathprint.cpp
File metadata and controls
86 lines (71 loc) · 2.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
#include <cstdint>
#include <cstring>
#include <sstream>
#include <nanobind/typing.h>
#include "mlx/utils.h"
#include "python/src/utils.h"
#include "mlx/mlx.h"
namespace mx = mlx::core;
namespace nb = nanobind;
using namespace nb::literals;
struct PrintOptionsContext {
mx::PrintOptions old_options;
mx::PrintOptions new_options;
PrintOptionsContext(mx::PrintOptions p) : new_options(p) {}
PrintOptionsContext& enter() {
old_options = mx::get_global_formatter().format_options;
mx::set_printoptions(new_options);
return *this;
}
void exit(nb::args) {
mx::set_printoptions(old_options);
}
};
void init_print(nb::module_& m) {
// Set Python print formatting options
mx::get_global_formatter().capitalize_bool = true;
// Expose printing options to Python: allow setting global precision.
nb::class_<mx::PrintOptions>(m, "PrintOptions")
.def(nb::init<int>(), "precision"_a = -1)
.def_rw("precision", &mx::PrintOptions::precision);
m.def(
"set_printoptions",
[](int precision) { mx::set_printoptions({precision}); },
"precision"_a = mx::get_global_formatter().format_options.precision,
R"pbdoc(
Set global printing precision for array formatting.
Example:
>>> print(x) # Uses default precision
>>> mx.set_printoptions(precision=3):
>>> print(x) # Uses precision of 3
>>> print(x) # Uses precision of 3 (again)
Args:
precision (int): Number of decimal places.
)pbdoc");
m.def(
"get_printoptions",
[]() { return mx::get_global_formatter().format_options; },
R"pbdoc(
Get global printing precision for array formatting.
Returns:
PrintOptions: The format options used for printing arrays.
)pbdoc");
nb::class_<PrintOptionsContext>(m, "_PrintOptionsContext")
.def(nb::init<mx::PrintOptions>())
.def("__enter__", &PrintOptionsContext::enter)
.def("__exit__", &PrintOptionsContext::exit);
m.def(
"printoptions",
[](int precision) { return PrintOptionsContext({precision}); },
"precision"_a = mx::get_global_formatter().format_options.precision,
R"pbdoc(
Context manager for setting print options temporarily.
Example:
>>> print(x) # Uses default precision
>>> with mx.printoptions(precision=3):
>>> print(x) # Uses precision of 3
>>> print(x) # Back to default precision
Args:
precision (int): Number of decimal places. Use -1 for default
)pbdoc");
}