diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 49643229..5f17f4b3 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -242,11 +242,13 @@ array linspace( if (num == 1) { return astype(array({start}), dtype, s); } - array t = divide(arange(0, num, float32, s), array(num - 1, float32), s); - array t_bar = subtract(array(1, float32), t, s); + auto inner_type = dtype == float64 ? float64 : float32; + array t = + divide(arange(0, num, inner_type, s), array(num - 1, inner_type), s); + array t_bar = subtract(array(1, inner_type), t, s); return astype( - add(multiply(t_bar, array(start, float32), s), - multiply(t, array(stop, float32), s), + add(multiply(t_bar, array(start, inner_type), s), + multiply(t, array(stop, inner_type), s), s), dtype, s); diff --git a/python/src/convert.cpp b/python/src/convert.cpp index 0367aceb..81adcaf1 100644 --- a/python/src/convert.cpp +++ b/python/src/convert.cpp @@ -259,6 +259,8 @@ nb::object tolist(mx::array& a) { return to_list(a, 0, 0); case mx::bfloat16: return to_list(a, 0, 0); + case mx::float64: + return to_list(a, 0, 0); case mx::complex64: return to_list>(a, 0, 0); default: diff --git a/python/tests/test_double.py b/python/tests/test_double.py index fccf3628..f957e8fc 100644 --- a/python/tests/test_double.py +++ b/python/tests/test_double.py @@ -292,6 +292,15 @@ class TestDouble(mlx_tests.MLXTestCase): b = np.array(a) self.assertTrue(np.array_equal(a, b)) + a = mx.array([1.0, 2.0], mx.float64) + b = a.tolist() + self.assertEqual(b, [1.0, 2.0]) + + def test_linspace(self): + with mx.stream(mx.cpu): + vals = mx.linspace(0, math.pi, 2, mx.float64) + self.assertEqual(vals.tolist()[1], math.pi) + if __name__ == "__main__": mlx_tests.MLXTestRunner()