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

Added formatter structure and a boolean value formatter (#354)

Browse Source 
* added formatter structure and a boolean value formatter

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
toji 2024-01-18 21:19:41 +05:30 committed by GitHub
parent d1fef34138
commit 49a52610b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 64 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
mlx/utils.cpp
View File

@ -7,6 +7,46 @@
namespace mlx::core {
void PrintFormatter::print(std::ostream& os, bool val) {
if (capitalize_bool) {
os << (val ? "True" : "False");
} else {
os << val;
}
}
inline void PrintFormatter::print(std::ostream& os, int16_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, uint16_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, int32_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, uint32_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, int64_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, uint64_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, float16_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, bfloat16_t val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, float val) {
os << val;
}
inline void PrintFormatter::print(std::ostream& os, complex64_t val) {
os << val;
}
PrintFormatter global_formatter;
Dtype result_type(const std::vector<array>& arrays) {
std::vector<Dtype> dtypes(1, bool_);
for (auto& arr : arrays) {
@ -136,7 +176,7 @@ void print_subarray(std::ostream& os, const array& a, size_t index, int dim) {
i = n - num_print - 1;
index += s * (n - 2 * num_print - 1);
} else if (is_last) {
os << a.data<T>()[index];
global_formatter.print(os, a.data<T>()[index]);
} else {
print_subarray<T>(os, a, index, dim + 1);
}
@ -153,7 +193,7 @@ void print_array(std::ostream& os, const array& a) {
os << "array(";
if (a.ndim() == 0) {
auto data = a.data<T>();
os << data[0];
global_formatter.print(os, data[0]);
} else {
print_subarray<T>(os, a, 0, 0);
}

18
mlx/utils.h
View File

@ -9,6 +9,24 @@
namespace mlx::core {
struct PrintFormatter {
inline void print(std::ostream& os, bool val);
inline void print(std::ostream& os, int16_t val);
inline void print(std::ostream& os, uint16_t val);
inline void print(std::ostream& os, int32_t val);
inline void print(std::ostream& os, uint32_t val);
inline void print(std::ostream& os, int64_t val);
inline void print(std::ostream& os, uint64_t val);
inline void print(std::ostream& os, float16_t val);
inline void print(std::ostream& os, bfloat16_t val);
inline void print(std::ostream& os, float val);
inline void print(std::ostream& os, complex64_t val);
bool capitalize_bool{false};
};
extern PrintFormatter global_formatter;
/** The type from promoting the arrays' types with one another. */
Dtype result_type(const std::vector<array>& arrays);

3
python/src/array.cpp
View File

@ -520,6 +520,9 @@ class ArrayPythonIterator {
};
void init_array(py::module_& m) {
// Set Python print formatting options
mlx::core::global_formatter.capitalize_bool = true;
// Types
py::class_<Dtype>(
m,

2
python/tests/test_array.py
View File

@ -304,7 +304,7 @@ class TestArray(mlx_tests.MLXTestCase):
def test_array_repr(self):
x = mx.array(True)
self.assertEqual(str(x), "array(true, dtype=bool)")
self.assertEqual(str(x), "array(True, dtype=bool)")
x = mx.array(1)
self.assertEqual(str(x), "array(1, dtype=int32)")
x = mx.array(1.0)