zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-16 13:51:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: xuan <xuan@apple.com>
This commit is contained in:
mutexuan 2024-01-02 11:02:04 +08:00 committed by GitHub
parent e09bf35b28
commit 350095ce6e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/src/array.cpp
View File

@ -64,7 +64,7 @@ auto to_scalar(array& a) {
case float32: case float32:
return py::cast(a.item<float>(retain_graph)); return py::cast(a.item<float>(retain_graph));
case bfloat16: case bfloat16:
return py::cast(static_cast<float>(a.item<float16_t>(retain_graph))); return py::cast(static_cast<float>(a.item<bfloat16_t>(retain_graph)));
case complex64: case complex64:
return py::cast(a.item<std::complex<float>>(retain_graph)); return py::cast(a.item<std::complex<float>>(retain_graph));
} }

3
python/tests/test_array.py
View File

@ -102,6 +102,9 @@ class TestArray(mlx_tests.MLXTestCase):
self.assertEqual(x.item(), 1) self.assertEqual(x.item(), 1)
self.assertTrue(isinstance(x.item(), int)) self.assertTrue(isinstance(x.item(), int))
x = mx.array(1, mx.bfloat16)
self.assertEqual(x.item(), 1.0)
x = mx.array(1.0) x = mx.array(1.0)
self.assertEqual(x.size, 1) self.assertEqual(x.size, 1)
self.assertEqual(x.ndim, 0) self.assertEqual(x.ndim, 0)