zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 18:28:12 +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
4 changed files with 64 additions and 3 deletions
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);