mlx 1.0.0

data/mlx/python/tests/test_device.py

@@ -0,0 +1,150 @@
+# Copyright © 2023 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+# Don't inherit from MLXTestCase to avoid call to setUp
+class TestDefaultDevice(unittest.TestCase):
+    def test_mlx_default_device(self):
+        device = mx.default_device()
+        if mx.is_available(mx.gpu):
+            self.assertEqual(device, mx.Device(mx.gpu))
+            self.assertEqual(str(device), "Device(gpu, 0)")
+            self.assertEqual(device, mx.gpu)
+            self.assertEqual(mx.gpu, device)
+        else:
+            self.assertEqual(device.type, mx.Device(mx.cpu))
+            with self.assertRaises(ValueError):
+                mx.set_default_device(mx.gpu)
+
+
+class TestDevice(mlx_tests.MLXTestCase):
+    def test_device(self):
+        device = mx.default_device()
+
+        cpu = mx.Device(mx.cpu)
+        mx.set_default_device(cpu)
+        self.assertEqual(mx.default_device(), cpu)
+        self.assertEqual(str(cpu), "Device(cpu, 0)")
+
+        mx.set_default_device(mx.cpu)
+        self.assertEqual(mx.default_device(), mx.cpu)
+        self.assertEqual(cpu, mx.cpu)
+        self.assertEqual(mx.cpu, cpu)
+
+        # Restore device
+        mx.set_default_device(device)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available")
+    def test_device_context(self):
+        default = mx.default_device()
+        diff = mx.cpu if default == mx.gpu else mx.gpu
+        self.assertNotEqual(default, diff)
+        with mx.stream(diff):
+            a = mx.add(mx.zeros((2, 2)), mx.ones((2, 2)))
+            mx.eval(a)
+            self.assertEqual(mx.default_device(), diff)
+        self.assertEqual(mx.default_device(), default)
+
+    def test_op_on_device(self):
+        x = mx.array(1.0)
+        y = mx.array(1.0)
+
+        a = mx.add(x, y, stream=None)
+        b = mx.add(x, y, stream=mx.default_device())
+        self.assertEqual(a.item(), b.item())
+        b = mx.add(x, y, stream=mx.cpu)
+        self.assertEqual(a.item(), b.item())
+
+        if mx.metal.is_available():
+            b = mx.add(x, y, stream=mx.gpu)
+            self.assertEqual(a.item(), b.item())
+
+
+class TestStream(mlx_tests.MLXTestCase):
+    def test_stream(self):
+        s1 = mx.default_stream(mx.default_device())
+        self.assertEqual(s1.device, mx.default_device())
+
+        s2 = mx.new_stream(mx.default_device())
+        self.assertEqual(s2.device, mx.default_device())
+        self.assertNotEqual(s1, s2)
+
+        if mx.is_available(mx.gpu):
+            s_gpu = mx.default_stream(mx.gpu)
+            self.assertEqual(s_gpu.device, mx.gpu)
+        else:
+            with self.assertRaises(ValueError):
+                mx.default_stream(mx.gpu)
+
+        s_cpu = mx.default_stream(mx.cpu)
+        self.assertEqual(s_cpu.device, mx.cpu)
+
+        s_cpu = mx.new_stream(mx.cpu)
+        self.assertEqual(s_cpu.device, mx.cpu)
+
+        if mx.is_available(mx.gpu):
+            s_gpu = mx.new_stream(mx.gpu)
+            self.assertEqual(s_gpu.device, mx.gpu)
+        else:
+            with self.assertRaises(ValueError):
+                mx.new_stream(mx.gpu)
+
+    def test_op_on_stream(self):
+        x = mx.array(1.0)
+        y = mx.array(1.0)
+
+        a = mx.add(x, y, stream=mx.default_stream(mx.default_device()))
+
+        if mx.is_available(mx.gpu):
+            b = mx.add(x, y, stream=mx.default_stream(mx.gpu))
+            self.assertEqual(a.item(), b.item())
+            s_gpu = mx.new_stream(mx.gpu)
+            b = mx.add(x, y, stream=s_gpu)
+            self.assertEqual(a.item(), b.item())
+
+        b = mx.add(x, y, stream=mx.default_stream(mx.cpu))
+        self.assertEqual(a.item(), b.item())
+        s_cpu = mx.new_stream(mx.cpu)
+        b = mx.add(x, y, stream=s_cpu)
+        self.assertEqual(a.item(), b.item())
+
+
+class TestDeviceInfo(mlx_tests.MLXTestCase):
+    def test_device_count(self):
+        cpu_count = mx.device_count(mx.cpu)
+        self.assertIsInstance(cpu_count, int)
+        self.assertEqual(cpu_count, 1)
+
+        gpu_count = mx.device_count(mx.gpu)
+        self.assertIsInstance(gpu_count, int)
+        self.assertGreaterEqual(gpu_count, 0)
+
+    def test_device_info_cpu(self):
+        info = mx.device_info(mx.cpu)
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+        self.assertTrue(len(info["device_name"]) > 0)
+        self.assertIn("architecture", info)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available")
+    def test_device_info_gpu(self):
+        gpu_count = mx.device_count(mx.gpu)
+        for i in range(gpu_count):
+            info = mx.device_info(mx.Device(mx.gpu, i))
+            self.assertIsInstance(info, dict)
+            self.assertIn("device_name", info)
+            self.assertTrue(len(info["device_name"]) > 0)
+            self.assertIn("architecture", info)
+
+    def test_device_info_default(self):
+        info = mx.device_info()
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner()

data/mlx/python/tests/test_double.py

@@ -0,0 +1,306 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+import os
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+import numpy as np
+
+
+class TestDouble(mlx_tests.MLXTestCase):
+    def test_unary_ops(self):
+        shape = (3, 3)
+        x = mx.random.normal(shape=shape)
+
+        if mx.default_device() == mx.gpu:
+            with self.assertRaises(ValueError):
+                x.astype(mx.float64)
+
+        x_double = x.astype(mx.float64, stream=mx.cpu)
+
+        ops = [
+            mx.abs,
+            mx.arccos,
+            mx.arccosh,
+            mx.arcsin,
+            mx.arcsinh,
+            mx.arctan,
+            mx.arctanh,
+            mx.ceil,
+            mx.erf,
+            mx.erfinv,
+            mx.exp,
+            mx.expm1,
+            mx.floor,
+            mx.log,
+            mx.logical_not,
+            mx.negative,
+            mx.round,
+            mx.sin,
+            mx.sinh,
+            mx.sqrt,
+            mx.rsqrt,
+            mx.tan,
+            mx.tanh,
+        ]
+        for op in ops:
+            if mx.default_device() == mx.gpu:
+                with self.assertRaises(ValueError):
+                    op(x_double)
+                continue
+            y = op(x)
+            y_double = op(x_double)
+            self.assertTrue(
+                mx.allclose(y, y_double.astype(mx.float32, mx.cpu), equal_nan=True)
+            )
+
+    def test_binary_ops(self):
+        shape = (3, 3)
+        a = mx.random.normal(shape=shape)
+        b = mx.random.normal(shape=shape)
+
+        a_double = a.astype(mx.float64, stream=mx.cpu)
+        b_double = b.astype(mx.float64, stream=mx.cpu)
+
+        ops = [
+            mx.add,
+            mx.arctan2,
+            mx.divide,
+            mx.multiply,
+            mx.subtract,
+            mx.logical_and,
+            mx.logical_or,
+            mx.remainder,
+            mx.maximum,
+            mx.minimum,
+            mx.power,
+            mx.equal,
+            mx.greater,
+            mx.greater_equal,
+            mx.less,
+            mx.less_equal,
+            mx.not_equal,
+            mx.logaddexp,
+        ]
+        for op in ops:
+            if mx.default_device() == mx.gpu:
+                with self.assertRaises(ValueError):
+                    op(a_double, b_double)
+                continue
+            y = op(a, b)
+            y_double = op(a_double, b_double)
+            self.assertTrue(
+                mx.allclose(y, y_double.astype(mx.float32, mx.cpu), equal_nan=True)
+            )
+
+    def test_where(self):
+        shape = (3, 3)
+        cond = mx.random.uniform(shape=shape) > 0.5
+        a = mx.random.normal(shape=shape)
+        b = mx.random.normal(shape=shape)
+
+        a_double = a.astype(mx.float64, stream=mx.cpu)
+        b_double = b.astype(mx.float64, stream=mx.cpu)
+
+        if mx.default_device() == mx.gpu:
+            with self.assertRaises(ValueError):
+                mx.where(cond, a_double, b_double)
+            return
+        y = mx.where(cond, a, b)
+        y_double = mx.where(cond, a_double, b_double)
+        self.assertTrue(mx.allclose(y, y_double.astype(mx.float32, mx.cpu)))
+
+    def test_reductions(self):
+        shape = (32, 32)
+        a = mx.random.normal(shape=shape)
+        a_double = a.astype(mx.float64, stream=mx.cpu)
+
+        axes = [0, 1, (0, 1)]
+        ops = [mx.sum, mx.prod, mx.min, mx.max, mx.any, mx.all]
+
+        for op in ops:
+            for ax in axes:
+                if mx.default_device() == mx.gpu:
+                    with self.assertRaises(ValueError):
+                        op(a_double, axis=ax)
+                    continue
+                y = op(a)
+                y_double = op(a_double)
+                self.assertTrue(mx.allclose(y, y_double.astype(mx.float32, mx.cpu)))
+
+    def test_get_and_set_item(self):
+        shape = (3, 3)
+        a = mx.random.normal(shape=shape)
+        b = mx.random.normal(shape=(2,))
+        a_double = a.astype(mx.float64, stream=mx.cpu)
+        b_double = b.astype(mx.float64, stream=mx.cpu)
+        idx_i = mx.array([0, 2])
+        idx_j = mx.array([0, 2])
+
+        if mx.default_device() == mx.gpu:
+            with self.assertRaises(ValueError):
+                a_double[idx_i, idx_j]
+        else:
+            y = a[idx_i, idx_j]
+            y_double = a_double[idx_i, idx_j]
+            self.assertTrue(mx.allclose(y, y_double.astype(mx.float32, mx.cpu)))
+
+        if mx.default_device() == mx.gpu:
+            with self.assertRaises(ValueError):
+                a_double[idx_i, idx_j] = b_double
+        else:
+            a[idx_i, idx_j] = b
+            a_double[idx_i, idx_j] = b_double
+            self.assertTrue(mx.allclose(a, a_double.astype(mx.float32, mx.cpu)))
+
+    def test_gemm(self):
+        shape = (8, 8)
+        a = mx.random.normal(shape=shape)
+        b = mx.random.normal(shape=shape)
+
+        a_double = a.astype(mx.float64, stream=mx.cpu)
+        b_double = b.astype(mx.float64, stream=mx.cpu)
+
+        if mx.default_device() == mx.gpu:
+            with self.assertRaises(ValueError):
+                a_double @ b_double
+            return
+        y = a @ b
+        y_double = a_double @ b_double
+        self.assertTrue(
+            mx.allclose(y, y_double.astype(mx.float32, mx.cpu), equal_nan=True)
+        )
+
+    def test_type_promotion(self):
+        import mlx.core as mx
+
+        a = mx.array([4, 8], mx.float64)
+        b = mx.array([4, 8], mx.int32)
+
+        with mx.stream(mx.cpu):
+            c = a + b
+            self.assertEqual(c.dtype, mx.float64)
+
+    def test_lapack(self):
+        with mx.stream(mx.cpu):
+            # QRF
+            A = mx.array([[2.0, 3.0], [1.0, 2.0]], dtype=mx.float64)
+            Q, R = mx.linalg.qr(A)
+            out = Q @ R
+            self.assertTrue(mx.allclose(out, A))
+            out = Q.T @ Q
+            self.assertTrue(mx.allclose(out, mx.eye(2)))
+            self.assertTrue(mx.allclose(mx.tril(R, -1), mx.zeros_like(R)))
+            self.assertEqual(Q.dtype, mx.float64)
+            self.assertEqual(R.dtype, mx.float64)
+
+            # SVD
+            A = mx.array(
+                [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float64
+            )
+            U, S, Vt = mx.linalg.svd(A)
+            self.assertTrue(mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, A))
+
+            # Inverse
+            A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float64)
+            A_inv = mx.linalg.inv(A)
+            self.assertTrue(mx.allclose(A @ A_inv, mx.eye(A.shape[0])))
+
+            # Tri inv
+            A = mx.array([[1, 0, 0], [6, -5, 0], [-9, 8, 7]], dtype=mx.float64)
+            B = mx.array([[7, 0, 0], [3, -2, 0], [1, 8, 3]], dtype=mx.float64)
+            AB = mx.stack([A, B])
+            invs = mx.linalg.tri_inv(AB, upper=False)
+            for M, M_inv in zip(AB, invs):
+                self.assertTrue(mx.allclose(M @ M_inv, mx.eye(M.shape[0])))
+
+            # Cholesky
+            sqrtA = mx.array(
+                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float64
+            )
+            A = sqrtA.T @ sqrtA / 81
+            L = mx.linalg.cholesky(A)
+            U = mx.linalg.cholesky(A, upper=True)
+            self.assertTrue(mx.allclose(L @ L.T, A))
+            self.assertTrue(mx.allclose(U.T @ U, A))
+
+            # Psueod inverse
+            A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float64)
+            A_plus = mx.linalg.pinv(A)
+            self.assertTrue(mx.allclose(A @ A_plus @ A, A))
+
+            # Eigh
+            def check_eigs_and_vecs(A_np, kwargs={}):
+                A = mx.array(A_np, dtype=mx.float64)
+                eig_vals, eig_vecs = mx.linalg.eigh(A, **kwargs)
+                eig_vals_np, _ = np.linalg.eigh(A_np, **kwargs)
+                self.assertTrue(np.allclose(eig_vals, eig_vals_np))
+                self.assertTrue(
+                    mx.allclose(A @ eig_vecs, eig_vals[..., None, :] * eig_vecs)
+                )
+
+                eig_vals_only = mx.linalg.eigvalsh(A, **kwargs)
+                self.assertTrue(mx.allclose(eig_vals, eig_vals_only))
+
+            # Test a simple 2x2 symmetric matrix
+            A_np = np.array([[1.0, 2.0], [2.0, 4.0]], dtype=np.float64)
+            check_eigs_and_vecs(A_np)
+
+            # Test a larger random symmetric matrix
+            n = 5
+            np.random.seed(1)
+            A_np = np.random.randn(n, n).astype(np.float64)
+            A_np = (A_np + A_np.T) / 2
+            check_eigs_and_vecs(A_np)
+
+            # Test with upper triangle
+            check_eigs_and_vecs(A_np, {"UPLO": "U"})
+
+            # LU factorization
+            # Test 3x3 matrix
+            a = mx.array(
+                [[3.0, 1.0, 2.0], [1.0, 8.0, 6.0], [9.0, 2.0, 5.0]], dtype=mx.float64
+            )
+            P, L, U = mx.linalg.lu(a)
+            self.assertTrue(mx.allclose(L[P, :] @ U, a))
+
+            # Solve triangular
+            # Test lower triangular matrix
+            a = mx.array(
+                [[4.0, 0.0, 0.0], [2.0, 3.0, 0.0], [1.0, -2.0, 5.0]], dtype=mx.float64
+            )
+            b = mx.array([8.0, 14.0, 3.0], dtype=mx.float64)
+
+            result = mx.linalg.solve_triangular(a, b, upper=False)
+            expected = np.linalg.solve(np.array(a), np.array(b))
+            self.assertTrue(np.allclose(result, expected))
+
+            # Test upper triangular matrix
+            a = mx.array(
+                [[3.0, 2.0, 1.0], [0.0, 5.0, 4.0], [0.0, 0.0, 6.0]], dtype=mx.float64
+            )
+            b = mx.array([13.0, 33.0, 18.0], dtype=mx.float64)
+
+            result = mx.linalg.solve_triangular(a, b, upper=True)
+            expected = np.linalg.solve(np.array(a), np.array(b))
+            self.assertTrue(np.allclose(result, expected))
+
+    def test_conversion(self):
+        a = mx.array([1.0, 2.0], mx.float64)
+        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()