[CUDA] ternary with select op (#2283)

* cuda ternary with select op

* comment + fix

* fix
This commit is contained in:
Awni Hannun
2025-06-12 20:24:43 -07:00
committed by GitHub
parent aa07429bad
commit 2188199ff8
5 changed files with 231 additions and 1 deletions

View File

@@ -0,0 +1,12 @@
// Copyright © 2025 Apple Inc.
namespace mlx::core::cu {
struct Select {
template <typename T>
__device__ T operator()(bool condition, T x, T y) {
return condition ? x : y;
}
};
} // namespace mlx::core::cu