fix type cast error in item() for bfloat16 (#339)

Co-authored-by: xuan <[email protected]>
2024-01-01 19:02:04 -08:00
parent e09bf35b28
commit 350095ce6e
2 changed files with 4 additions and 1 deletions
@@ -102,6 +102,9 @@ class TestArray(mlx_tests.MLXTestCase):
        self.assertEqual(x.item(), 1)
        self.assertTrue(isinstance(x.item(), int))

+        x = mx.array(1, mx.bfloat16)
+        self.assertEqual(x.item(), 1.0)
+
        x = mx.array(1.0)
        self.assertEqual(x.size, 1)
        self.assertEqual(x.ndim, 0)