Fix assigning bool to float16/bfloat16 (#3229)

Co-authored-by: KD2YCU <me@kd2ycu.com>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

Author: 3 people

mlx/types/bf16.h

@@ -8,6 +8,7 @@
 #include <vector>
 
 #define __MLX_BFLOAT_NAN__ 0x7FC0
+#define __MLX_BFLOAT_ONE__ 0x3F80
 
 namespace mlx::core {
 
@@ -29,8 +30,8 @@ struct _MLX_BFloat16 {
 
   // Appease std::vector<bool> for being special
   _MLX_BFloat16& operator=(std::vector<bool>::reference x) {
-    bits_ = x;
-    return *this;
+    bits_ = (x) ? __MLX_BFLOAT_ONE__ : 0;
+    return (*this);
   }
 
   _MLX_BFloat16& operator=(const float& x) {

mlx/types/fp16.h

@@ -8,6 +8,7 @@
 #include <vector>
 
 #define __MLX_HALF_NAN__ 0x7D00
+#define __MLX_HALF_ONE__ 0x3C00
 
 namespace mlx::core {
 
@@ -29,8 +30,8 @@ struct _MLX_Float16 {
 
   // Appease std::vector<bool> for being special
   _MLX_Float16& operator=(std::vector<bool>::reference x) {
-    bits_ = x;
-    return *this;
+    bits_ = (x) ? __MLX_HALF_ONE__ : 0;
+    return (*this);
   }
 
   _MLX_Float16& operator=(const float& x) {

tests/array_tests.cpp

@@ -105,6 +105,16 @@ TEST_CASE("test array basics") {
     CHECK_EQ(x.dtype(), bool_);
     CHECK(array_equal(x, array({false, true, false, true})).item<bool>());
   }
+
+  // Regression: vector<bool>::reference to fp16/bf16 stored raw bits
+  {
+    std::vector<bool> data = {true, false, true};
+    auto bf = array(data.begin(), {3}, bfloat16);
+    CHECK(array_equal(bf, array({1.0f, 0.0f, 1.0f}, bfloat16)).item<bool>());
+
+    auto fp = array(data.begin(), {3}, float16);
+    CHECK(array_equal(fp, array({1.0f, 0.0f, 1.0f}, float16)).item<bool>());
+  }
 }
 
 TEST_CASE("test array types") {