diff --git a/mlx/types/bf16.h b/mlx/types/bf16.h index 59519417..7feaaa66 100644 --- a/mlx/types/bf16.h +++ b/mlx/types/bf16.h @@ -8,6 +8,7 @@ #include #define __MLX_BFLOAT_NAN__ 0x7FC0 +#define __MLX_BFLOAT_ONE__ 0x3F80 namespace mlx::core { @@ -29,8 +30,8 @@ struct _MLX_BFloat16 { // Appease std::vector for being special _MLX_BFloat16& operator=(std::vector::reference x) { - bits_ = x; - return *this; + bits_ = (x) ? __MLX_BFLOAT_ONE__ : 0; + return (*this); } _MLX_BFloat16& operator=(const float& x) { diff --git a/mlx/types/fp16.h b/mlx/types/fp16.h index c174afeb..31b0a78d 100644 --- a/mlx/types/fp16.h +++ b/mlx/types/fp16.h @@ -8,6 +8,7 @@ #include #define __MLX_HALF_NAN__ 0x7D00 +#define __MLX_HALF_ONE__ 0x3C00 namespace mlx::core { @@ -29,8 +30,8 @@ struct _MLX_Float16 { // Appease std::vector for being special _MLX_Float16& operator=(std::vector::reference x) { - bits_ = x; - return *this; + bits_ = (x) ? __MLX_HALF_ONE__ : 0; + return (*this); } _MLX_Float16& operator=(const float& x) { diff --git a/tests/array_tests.cpp b/tests/array_tests.cpp index 84909472..6e0d0346 100644 --- a/tests/array_tests.cpp +++ b/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()); } + + // Regression: vector::reference to fp16/bf16 stored raw bits + { + std::vector data = {true, false, true}; + auto bf = array(data.begin(), {3}, bfloat16); + CHECK(array_equal(bf, array({1.0f, 0.0f, 1.0f}, bfloat16)).item()); + + auto fp = array(data.begin(), {3}, float16); + CHECK(array_equal(fp, array({1.0f, 0.0f, 1.0f}, float16)).item()); + } } TEST_CASE("test array types") {