mlx 1.0.0

data/mlx/python/tests/test_random.py

@@ -0,0 +1,392 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+class TestRandom(mlx_tests.MLXTestCase):
+    def test_global_rng(self):
+        mx.random.seed(3)
+        a = mx.random.uniform()
+        b = mx.random.uniform()
+
+        mx.random.seed(3)
+        x = mx.random.uniform()
+        y = mx.random.uniform()
+
+        self.assertEqual(a.item(), x.item())
+        self.assertEqual(y.item(), b.item())
+
+    def test_key(self):
+        k1 = mx.random.key(0)
+        k2 = mx.random.key(0)
+        self.assertTrue(mx.array_equal(k1, k2))
+
+        k2 = mx.random.key(1)
+        self.assertFalse(mx.array_equal(k1, k2))
+
+    def test_key_split(self):
+        key = mx.random.key(0)
+
+        k1, k2 = mx.random.split(key)
+        self.assertFalse(mx.array_equal(k1, k2))
+
+        r1, r2 = mx.random.split(key)
+        self.assertTrue(mx.array_equal(k1, r1))
+        self.assertTrue(mx.array_equal(k2, r2))
+
+        keys = mx.random.split(key, 10)
+        self.assertEqual(keys.shape, (10, 2))
+
+    def test_uniform(self):
+        key = mx.random.key(0)
+        a = mx.random.uniform(key=key)
+        self.assertEqual(a.shape, ())
+        self.assertEqual(a.dtype, mx.float32)
+
+        b = mx.random.uniform(key=key)
+        self.assertEqual(a.item(), b.item())
+
+        a = mx.random.uniform(shape=(2, 3))
+        self.assertEqual(a.shape, (2, 3))
+
+        a = mx.random.uniform(shape=(1000,), low=-1, high=5)
+        self.assertTrue(mx.all((a > -1) < 5).item())
+
+        a = mx.random.uniform(shape=(1000,), low=mx.array(-1), high=5)
+        self.assertTrue(mx.all((a > -1) < 5).item())
+
+        a = mx.random.uniform(low=-0.1, high=0.1, shape=(1,), dtype=mx.bfloat16)
+        self.assertEqual(a.dtype, mx.bfloat16)
+
+        self.assertEqual(mx.random.uniform().dtype, mx.random.uniform(dtype=None).dtype)
+
+    def test_normal_and_laplace(self):
+        # Same tests for normal and laplace.
+        for distribution_sampler in [mx.random.normal, mx.random.laplace]:
+            key = mx.random.key(0)
+            a = distribution_sampler(key=key)
+            self.assertEqual(a.shape, ())
+            self.assertEqual(a.dtype, mx.float32)
+
+            b = distribution_sampler(key=key)
+            self.assertEqual(a.item(), b.item())
+
+            a = distribution_sampler(shape=(2, 3))
+            self.assertEqual(a.shape, (2, 3))
+
+            ## Generate in float16 or bfloat16
+            for t in [mx.float16, mx.bfloat16]:
+                a = distribution_sampler(dtype=t)
+                self.assertEqual(a.dtype, t)
+
+            # Generate with a given mean and standard deviation
+            loc = 1.0
+            scale = 2.0
+
+            a = distribution_sampler(shape=(3, 2), loc=loc, scale=scale, key=key)
+            b = scale * distribution_sampler(shape=(3, 2), key=key) + loc
+            self.assertTrue(mx.allclose(a, b))
+
+            a = distribution_sampler(
+                shape=(3, 2), loc=loc, scale=scale, dtype=mx.float16, key=key
+            )
+            b = (
+                scale * distribution_sampler(shape=(3, 2), dtype=mx.float16, key=key)
+                + loc
+            )
+            self.assertTrue(mx.allclose(a, b))
+
+            self.assertEqual(
+                distribution_sampler().dtype, distribution_sampler(dtype=None).dtype
+            )
+
+            # Test not getting -inf or inf with half precison
+            for hp in [mx.float16, mx.bfloat16]:
+                a = abs(distribution_sampler(shape=(10000,), loc=0, scale=1, dtype=hp))
+                self.assertTrue(mx.all(a < mx.inf))
+
+    def test_multivariate_normal(self):
+        key = mx.random.key(0)
+        mean = mx.array([0, 0])
+        cov = mx.array([[1, 0], [0, 1]])
+
+        a = mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)
+        self.assertEqual(a.shape, (2,))
+
+        ## Check dtypes
+        for t in [mx.float32]:
+            a = mx.random.multivariate_normal(
+                mean, cov, dtype=t, key=key, stream=mx.cpu
+            )
+            self.assertEqual(a.dtype, t)
+        for t in [
+            mx.int8,
+            mx.int32,
+            mx.int64,
+            mx.uint8,
+            mx.uint32,
+            mx.uint64,
+            mx.float16,
+            mx.bfloat16,
+        ]:
+            with self.assertRaises(ValueError):
+                mx.random.multivariate_normal(
+                    mean, cov, dtype=t, key=key, stream=mx.cpu
+                )
+
+        ## Check incompatible shapes
+        with self.assertRaises(ValueError):
+            mean = mx.zeros((2, 2))
+            cov = mx.zeros((2, 2))
+            mx.random.multivariate_normal(mean, cov, shape=(3,), key=key, stream=mx.cpu)
+
+        with self.assertRaises(ValueError):
+            mean = mx.zeros((2))
+            cov = mx.zeros((2, 2, 2))
+            mx.random.multivariate_normal(mean, cov, shape=(3,), key=key, stream=mx.cpu)
+
+        with self.assertRaises(ValueError):
+            mean = mx.zeros((3,))
+            cov = mx.zeros((2, 2))
+            mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)
+
+        with self.assertRaises(ValueError):
+            mean = mx.zeros((2,))
+            cov = mx.zeros((2, 3))
+            mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)
+
+        ## Different shape of mean and cov
+        mean = mx.array([[0, 7], [1, 2], [3, 4]])
+        cov = mx.array([[1, 0.5], [0.5, 1]])
+        a = mx.random.multivariate_normal(mean, cov, shape=(4, 3), stream=mx.cpu)
+        self.assertEqual(a.shape, (4, 3, 2))
+
+        ## Check correcteness of the mean and covariance
+        n_test = int(1e5)
+
+        def check_jointly_gaussian(data, mean, cov):
+            empirical_mean = mx.mean(data, axis=0)
+            empirical_cov = (
+                (data - empirical_mean).T @ (data - empirical_mean) / data.shape[0]
+            )
+            N = data.shape[1]
+            self.assertTrue(
+                mx.allclose(
+                    empirical_mean, mean, rtol=0.0, atol=10 * N**2 / math.sqrt(n_test)
+                )
+            )
+            self.assertTrue(
+                mx.allclose(
+                    empirical_cov, cov, rtol=0.0, atol=10 * N**2 / math.sqrt(n_test)
+                )
+            )
+
+        mean = mx.array([4.0, 7.0])
+        cov = mx.array([[2, 0.5], [0.5, 1]])
+        data = mx.random.multivariate_normal(
+            mean, cov, shape=(n_test,), key=key, stream=mx.cpu
+        )
+        check_jointly_gaussian(data, mean, cov)
+
+        mean = mx.arange(3)
+        cov = mx.array([[1, -1, 0.5], [-1, 1, -0.5], [0.5, -0.5, 1]])
+        data = mx.random.multivariate_normal(
+            mean, cov, shape=(n_test,), key=key, stream=mx.cpu
+        )
+        check_jointly_gaussian(data, mean, cov)
+
+    def test_randint(self):
+        a = mx.random.randint(0, 1, [])
+        self.assertEqual(a.shape, ())
+        self.assertEqual(a.dtype, mx.int32)
+
+        shape = (88,)
+        low = mx.array(3)
+        high = mx.array(15)
+
+        key = mx.random.key(0)
+        a = mx.random.randint(low, high, shape, key=key)
+        self.assertEqual(a.shape, shape)
+        self.assertEqual(a.dtype, mx.int32)
+
+        # Check using the same key yields the same value
+        b = mx.random.randint(low, high, shape, key=key)
+        self.assertListEqual(a.tolist(), b.tolist())
+
+        shape = (3, 4)
+        low = mx.reshape(mx.array([0] * 3), [3, 1])
+        high = mx.reshape(mx.array([12, 13, 14, 15]), [1, 4])
+
+        a = mx.random.randint(low, high, shape)
+        self.assertEqual(a.shape, shape)
+
+        a = mx.random.randint(-10, 10, [1000, 1000])
+        self.assertTrue(mx.all(-10 <= a).item() and mx.all(a < 10).item())
+
+        a = mx.random.randint(10, -10, [1000, 1000])
+        self.assertTrue(mx.all(a == 10).item())
+
+        self.assertEqual(
+            mx.random.randint(0, 1).dtype, mx.random.randint(0, 1, dtype=None).dtype
+        )
+
+    def test_bernoulli(self):
+        a = mx.random.bernoulli()
+        self.assertEqual(a.shape, ())
+        self.assertEqual(a.dtype, mx.bool_)
+
+        a = mx.random.bernoulli(mx.array(0.5), [5])
+        self.assertEqual(a.shape, (5,))
+
+        a = mx.random.bernoulli(mx.array([2.0, -2.0]))
+        self.assertEqual(a.tolist(), [True, False])
+        self.assertEqual(a.shape, (2,))
+
+        p = mx.array([0.1, 0.2, 0.3])
+        mx.reshape(p, [1, 3])
+        x = mx.random.bernoulli(p, [4, 3])
+        self.assertEqual(x.shape, (4, 3))
+
+        with self.assertRaises(ValueError):
+            mx.random.bernoulli(p, [2])  # Bad shape
+
+        with self.assertRaises(ValueError):
+            mx.random.bernoulli(0, [2])  # Bad type
+
+    def test_truncated_normal(self):
+        a = mx.random.truncated_normal(-2.0, 2.0)
+        self.assertEqual(a.size, 1)
+        self.assertEqual(a.dtype, mx.float32)
+
+        a = mx.random.truncated_normal(mx.array([]), mx.array([]))
+        self.assertEqual(a.dtype, mx.float32)
+        self.assertEqual(a.size, 0)
+
+        lower = mx.reshape(mx.array([-2.0, 0.0]), [1, 2])
+        upper = mx.reshape(mx.array([0.0, 1.0, 2.0]), [3, 1])
+        a = mx.random.truncated_normal(lower, upper)
+
+        self.assertEqual(a.shape, (3, 2))
+        self.assertTrue(mx.all(lower <= a).item() and mx.all(a <= upper).item())
+
+        a = mx.random.truncated_normal(2.0, -2.0)
+        self.assertTrue(mx.all(a == 2.0).item())
+
+        a = mx.random.truncated_normal(-3.0, 3.0, [542, 399])
+        self.assertEqual(a.shape, (542, 399))
+
+        lower = mx.array([-2.0, -1.0])
+        higher = mx.array([1.0, 2.0, 3.0])
+        with self.assertRaises(ValueError):
+            mx.random.truncated_normal(lower, higher)  # Bad shape
+
+        self.assertEqual(
+            mx.random.truncated_normal(0, 1).dtype,
+            mx.random.truncated_normal(0, 1, dtype=None).dtype,
+        )
+
+    def test_gumbel(self):
+        samples = mx.random.gumbel(shape=(100, 100))
+        self.assertEqual(samples.shape, (100, 100))
+        self.assertEqual(samples.dtype, mx.float32)
+        mean = 0.5772
+        # Std deviation of the sample mean is small (<0.02),
+        # so this test is pretty conservative
+        self.assertTrue(mx.abs(mx.mean(samples) - mean) < 0.2)
+
+        self.assertEqual(
+            mx.random.gumbel((1, 1)).dtype, mx.random.gumbel((1, 1), dtype=None).dtype
+        )
+
+    def test_categorical(self):
+        logits = mx.zeros((10, 20))
+        self.assertEqual(mx.random.categorical(logits, -1).shape, (10,))
+        self.assertEqual(mx.random.categorical(logits, 0).shape, (20,))
+        self.assertEqual(mx.random.categorical(logits, 1).shape, (10,))
+
+        out = mx.random.categorical(logits)
+        self.assertEqual(out.shape, (10,))
+        self.assertEqual(out.dtype, mx.uint32)
+        self.assertTrue(mx.max(out).item() < 20)
+
+        out = mx.random.categorical(logits, 0, [5, 20])
+        self.assertEqual(out.shape, (5, 20))
+        self.assertTrue(mx.max(out).item() < 10)
+
+        out = mx.random.categorical(logits, 1, num_samples=7)
+        self.assertEqual(out.shape, (10, 7))
+        out = mx.random.categorical(logits, 0, num_samples=7)
+        self.assertEqual(out.shape, (20, 7))
+
+        with self.assertRaises(ValueError):
+            mx.random.categorical(logits, shape=[10, 5], num_samples=5)
+
+    def test_permutation(self):
+        x = sorted(mx.random.permutation(4).tolist())
+        self.assertEqual([0, 1, 2, 3], x)
+
+        x = mx.array([0, 1, 2, 3])
+        x = sorted(mx.random.permutation(x).tolist())
+        self.assertEqual([0, 1, 2, 3], x)
+
+        x = mx.array([0, 1, 2, 3])
+        x = sorted(mx.random.permutation(x).tolist())
+
+        # 2-D
+        x = mx.arange(16).reshape(4, 4)
+        out = mx.sort(mx.random.permutation(x, axis=0), axis=0)
+        self.assertTrue(mx.array_equal(x, out))
+        out = mx.sort(mx.random.permutation(x, axis=1), axis=1)
+        self.assertTrue(mx.array_equal(x, out))
+
+        # Basically 0 probability this should fail.
+        sorted_x = mx.arange(16384)
+        x = mx.random.permutation(16384)
+        self.assertFalse(mx.array_equal(sorted_x, x))
+
+        # Preserves shape / doesn't cast input to int
+        x = mx.random.permutation(mx.array([[1]]))
+        self.assertEqual(x.shape, (1, 1))
+
+    def test_complex_normal(self):
+        sample = mx.random.normal(tuple(), dtype=mx.complex64)
+        self.assertEqual(sample.shape, tuple())
+        self.assertEqual(sample.dtype, mx.complex64)
+
+        sample = mx.random.normal((1, 2, 3, 4), dtype=mx.complex64)
+        self.assertEqual(sample.shape, (1, 2, 3, 4))
+        self.assertEqual(sample.dtype, mx.complex64)
+
+        sample = mx.random.normal((1, 2, 3, 4), dtype=mx.complex64, scale=2.0, loc=3.0)
+        self.assertEqual(sample.shape, (1, 2, 3, 4))
+        self.assertEqual(sample.dtype, mx.complex64)
+
+        sample = mx.random.normal(
+            (1, 2, 3, 4), dtype=mx.complex64, scale=2.0, loc=3.0 + 1j
+        )
+        self.assertEqual(sample.shape, (1, 2, 3, 4))
+        self.assertEqual(sample.dtype, mx.complex64)
+
+    def test_broadcastable_scale_loc(self):
+        b = mx.random.normal((10, 2))
+        sample = mx.random.normal((2, 10, 2), loc=b, scale=b)
+        mx.eval(sample)
+        self.assertEqual(sample.shape, (2, 10, 2))
+
+        with self.assertRaises(ValueError):
+            b = mx.random.normal((10,))
+            sample = mx.random.normal((2, 10, 2), loc=b, scale=b)
+
+        b = mx.random.normal((3, 1, 2))
+        sample = mx.random.normal((3, 4, 2), dtype=mx.float16, loc=b, scale=b)
+        mx.eval(sample)
+        self.assertEqual(sample.shape, (3, 4, 2))
+        self.assertEqual(sample.dtype, mx.float16)
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner()

data/mlx/python/tests/test_reduce.py

@@ -0,0 +1,223 @@
+# Copyright © 2023 Apple Inc.
+
+import unittest
+from itertools import combinations, permutations
+
+import mlx.core as mx
+import mlx_tests
+import numpy as np
+
+
+class TestReduce(mlx_tests.MLXTestCase):
+    def test_axis_permutation_sums(self):
+        for shape in [(5, 5, 1, 5, 5), (65, 65, 1, 65)]:
+            with self.subTest(shape=shape):
+                x_npy = (np.random.randn(*shape) * 128).astype(np.int32)
+                x_mlx = mx.array(x_npy)
+                for t in permutations(range(len(shape))):
+                    with self.subTest(t=t):
+                        y_npy = np.transpose(x_npy, t)
+                        y_mlx = mx.transpose(x_mlx, t)
+                        for n in range(1, len(shape) + 1):
+                            for a in combinations(range(len(shape)), n):
+                                with self.subTest(a=a):
+                                    z_npy = np.sum(y_npy, axis=a)
+                                    z_mlx = mx.sum(y_mlx, axis=a)
+                                    mx.eval(z_mlx)
+                                    self.assertTrue(np.all(z_npy == z_mlx))
+
+    def test_expand_sums(self):
+        x_npy = np.random.randn(5, 1, 5, 1, 5, 1).astype(np.float32)
+        x_mlx = mx.array(x_npy)
+        for m in range(1, 4):
+            for ax in combinations([1, 3, 5], m):
+                shape = np.array([5, 1, 5, 1, 5, 1])
+                shape[list(ax)] = 5
+                shape = shape.tolist()
+                with self.subTest(shape=shape):
+                    y_npy = np.broadcast_to(x_npy, shape)
+                    y_mlx = mx.broadcast_to(x_mlx, shape)
+                    for n in range(1, 7):
+                        for a in combinations(range(6), n):
+                            with self.subTest(a=a):
+                                z_npy = np.sum(y_npy, axis=a) / 1000
+                                z_mlx = mx.sum(y_mlx, axis=a) / 1000
+                                mx.eval(z_mlx)
+                                self.assertTrue(
+                                    np.allclose(z_npy, np.array(z_mlx), atol=1e-4)
+                                )
+
+    def test_dtypes(self):
+        int_dtypes = [
+            "int8",
+            "int16",
+            "int32",
+            "uint8",
+            "uint16",
+            "uint32",
+            "int64",
+            "uint64",
+            "complex64",
+        ]
+        float_dtypes = ["float32"]
+
+        for dtype in int_dtypes + float_dtypes:
+            with self.subTest(dtype=dtype):
+                x = np.random.uniform(0, 2, size=(3, 3, 3)).astype(getattr(np, dtype))
+                y = mx.array(x)
+
+                for op in ("sum", "prod", "min", "max"):
+                    with self.subTest(op=op):
+                        np_op = getattr(np, op)
+                        mlx_op = getattr(mx, op)
+
+                        for axes in (None, 0, 1, 2, (0, 1), (0, 2), (1, 2), (0, 1, 2)):
+                            with self.subTest(axes=axes):
+                                if op in ("sum", "prod"):
+                                    r_np = np_op(
+                                        x, axis=axes, dtype=(getattr(np, dtype))
+                                    )
+                                else:
+                                    r_np = np_op(x, axis=axes)
+                                r_mlx = mlx_op(y, axis=axes)
+                                mx.eval(r_mlx)
+                                self.assertTrue(np.allclose(r_np, r_mlx, atol=1e-4))
+
+    def test_arg_reduce(self):
+        dtypes = [
+            "uint8",
+            "uint16",
+            "uint32",
+            "uint64",
+            "int8",
+            "int16",
+            "int32",
+            "int64",
+            "float16",
+            "float32",
+        ]
+        for dtype in dtypes:
+            with self.subTest(dtype=dtype):
+                data = np.random.rand(10, 12, 13).astype(getattr(np, dtype))
+                x = mx.array(data)
+                for op in ["argmin", "argmax"]:
+                    for axis in range(3):
+                        for kd in [True, False]:
+                            a = getattr(mx, op)(x, axis, kd)
+                            b = getattr(np, op)(data, axis, keepdims=kd)
+                            self.assertEqual(a.tolist(), b.tolist())
+
+                for op in ["argmin", "argmax"]:
+                    a = getattr(mx, op)(x, keepdims=True)
+                    b = getattr(np, op)(data, keepdims=True)
+                    self.assertEqual(a.tolist(), b.tolist())
+                    a = getattr(mx, op)(x)
+                    b = getattr(np, op)(data)
+                    self.assertEqual(a.item(), b)
+
+    def test_edge_case(self):
+        x = (mx.random.normal((100, 1, 100, 100)) * 128).astype(mx.int32)
+        x = x.transpose(0, 3, 1, 2)
+
+        y = x.sum((0, 2, 3))
+        mx.eval(y)
+        z = np.array(x).sum((0, 2, 3))
+        self.assertTrue(np.all(z == y))
+
+    def test_sum_bool(self):
+        x = np.random.uniform(0, 1, size=(10, 10, 10)) > 0.5
+        y = mx.array(x)
+        npsum = x.sum().item()
+        mxsum = y.sum().item()
+        self.assertEqual(npsum, mxsum)
+
+    def test_many_reduction_axes(self):
+
+        def check(x, axes):
+            expected = x
+            for ax in axes:
+                expected = mx.sum(expected, axis=ax, keepdims=True)
+            out = mx.sum(x, axis=axes, keepdims=True)
+            self.assertTrue(mx.array_equal(out, expected))
+
+        x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4))
+        check(x, (0, 2, 4))
+
+        x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4))
+        check(x, (0, 2, 4, 6))
+
+        x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4))
+        check(x, (0, 2, 4, 6, 8))
+
+        x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4, 128))
+        x = x.transpose(1, 0, 2, 3, 4, 5, 6, 7, 8, 9)
+        check(x, (1, 3, 5, 7, 9))
+
+    def test_nan_propagation(self):
+        dtypes = [
+            "uint8",
+            "uint16",
+            "uint32",
+            "int8",
+            "int16",
+            "int32",
+            "float16",
+            "float32",
+        ]
+
+        for dtype in dtypes:
+            with self.subTest(dtype=dtype):
+                x = (mx.random.normal((4, 4)) * 10).astype(getattr(mx, dtype))
+                indices = mx.random.randint(0, 4, shape=(6,)).reshape(3, 2)
+                for idx in indices:
+                    x[idx[0], idx[1]] = mx.nan
+                x_np = np.array(x)
+
+                for op in ["max", "min"]:
+                    for axis in [0, 1]:
+                        out = getattr(mx, op)(x, axis=axis)
+                        ref = getattr(np, op)(x_np, axis=axis)
+                        self.assertTrue(np.array_equal(out, ref, equal_nan=True))
+
+    def test_nan_propagation_complex64(self):
+        complex_array_1 = mx.array(
+            [1 + 1j, 2 + 2j, 3 + 3j, mx.nan + 4j], dtype=mx.complex64
+        ).reshape(2, 2)
+        complex_array_2 = mx.array(
+            [1 + 1j, 2 + 2j, 3 + mx.nan * 1j, 4 + 4j], dtype=mx.complex64
+        ).reshape(2, 2)
+        complex_array_3 = mx.array(
+            [1 + 1j, 2 + mx.nan * 1j, 3 + 3j, 4 + 4j], dtype=mx.complex64
+        ).reshape(2, 2)
+        complex_array_4 = mx.array(
+            [mx.nan + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=mx.complex64
+        ).reshape(2, 2)
+
+        np_arrays = [
+            np.array(complex_array_1),
+            np.array(complex_array_2),
+            np.array(complex_array_3),
+            np.array(complex_array_4),
+        ]
+
+        for mx_arr, np_arr in zip(
+            [complex_array_1, complex_array_2, complex_array_3, complex_array_4],
+            np_arrays,
+        ):
+            for axis in [0, 1]:
+                for op in ["max", "min"]:
+                    out = getattr(mx, op)(mx_arr, axis=axis)
+                    ref = getattr(np, op)(np_arr, axis=axis)
+                    self.assertTrue(np.array_equal(out, ref, equal_nan=True))
+
+    def test_long_column(self):
+        a = (np.random.randn(8192, 64) * 32).astype(np.int32)
+        b = mx.array(a)
+
+        c1 = a.sum(0)
+        c2 = b.sum(0)
+        self.assertTrue(np.all(c1 == c2))
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner(failfast=True)