2292557fd6
dist/include/ contains the MLX and MLX-C headers needed for CGo compilation. Without these, go-mlx cannot be used as a module dependency (headers not found in module cache). Libraries (dylib/metallib) are still gitignored — users build those locally via cmake. Co-Authored-By: Virgil <virgil@lethean.io>
38 lines
735 B
C++
38 lines
735 B
C++
// Copyright © 2025 Apple Inc.
|
|
|
|
namespace mlx::core::distributed::detail {
|
|
|
|
template <typename T>
|
|
struct SumOp {
|
|
void operator()(const T* input, T* output, size_t N) const {
|
|
while (N-- > 0) {
|
|
*output += *input;
|
|
input++;
|
|
output++;
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct MaxOp {
|
|
void operator()(const T* input, T* output, size_t N) const {
|
|
while (N-- > 0) {
|
|
*output = std::max(*output, *input);
|
|
input++;
|
|
output++;
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct MinOp {
|
|
void operator()(const T* input, T* output, size_t N) const {
|
|
while (N-- > 0) {
|
|
*output = std::min(*output, *input);
|
|
input++;
|
|
output++;
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace mlx::core::distributed::detail
|