diff --git a/mlx/backend/metal/compiled.cpp b/mlx/backend/metal/compiled.cpp index cda06143d6..e95d7b8d5f 100644 --- a/mlx/backend/metal/compiled.cpp +++ b/mlx/backend/metal/compiled.cpp @@ -208,7 +208,7 @@ inline void build_kernel( " {0} tmp_{1} = ", get_type_string(x.dtype()), namer.get_name(x)); if (is_static_cast(x.primitive())) { os += fmt::format( - "static_cast<{0}>(tmp_{1});\n", + "cast_to<{0}>(tmp_{1});\n", get_type_string(x.dtype()), namer.get_name(x.inputs()[0])); } else { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 5eaa6cb955..663ec295b0 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1623,6 +1623,16 @@ def test_compile_abs_unsigned(self): x = mx.array([1, 2, 3], dtype) self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + def test_compiled_subnormal_bool_cast(self): + f32_sub = mx.array(np.array([0x00000001] * 4, dtype=np.uint32)).view(mx.float32) + f16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.float16) + bf16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.bfloat16) + + # A single-op compile does not fuse; the fused path needs >= 2 ops. + fn = mx.compile(lambda x: mx.broadcast_to(x, (2, 4)).astype(mx.bool_)) + for sub in (f32_sub, f16_sub, bf16_sub): + self.assertTrue(mx.all(fn(sub)).item()) + def test_compile_different_log_bases(self): # The logs are intermediates, since outputs are not simplified. def entropies(p):