mlx 1.0.0

data/mlx/python/tests/test_einsum.py

@@ -0,0 +1,363 @@
+# Copyright © 2024 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+import numpy as np
+
+
+class TestEinsum(mlx_tests.MLXTestCase):
+
+    def test_simple_path(self):
+        a = mx.zeros((5, 5))
+        path = mx.einsum_path("ii", a)
+        self.assertEqual(path[0], [(0,)])
+
+        path = mx.einsum_path("ij->i", a)
+        self.assertEqual(path[0], [(0,)])
+
+        path = mx.einsum_path("ii->i", a)
+        self.assertEqual(path[0], [(0,)])
+
+        a = mx.zeros((5, 8))
+        b = mx.zeros((8, 3))
+        path = mx.einsum_path("ij,jk", a, b)
+        self.assertEqual(path[0], [(0, 1)])
+        path = mx.einsum_path("ij,jk -> ijk", a, b)
+        self.assertEqual(path[0], [(0, 1)])
+
+        a = mx.zeros((5, 8))
+        b = mx.zeros((8, 3))
+        c = mx.zeros((3, 7))
+        path = mx.einsum_path("ij,jk,kl", a, b, c)
+
+        self.assertEqual(path[0], [(0, 1), (0, 1)])
+
+        a = mx.zeros((5, 8))
+        b = mx.zeros((8, 10))
+        c = mx.zeros((10, 7))
+        path = mx.einsum_path("ij,jk,kl", a, b, c)
+        self.assertEqual(path[0], [(1, 2), (0, 1)])
+
+    def test_longer_paths(self):
+        chars = "abcdefghijklmopqABC"
+        sizes = [2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3, 2, 3, 4]
+        dim_dict = {c: s for c, s in zip(chars, sizes)}
+        cases = [
+            "eb,cb,fb->cef",
+            "dd,fb,be,cdb->cef",
+            "dd,fb,be,cdb->cef",
+            "bca,cdb,dbf,afc->",
+            "dcc,fce,ea,dbf->ab",
+            "dcc,fce,ea,dbf->ab",
+        ]
+
+        for case in cases:
+            subscripts = case[: case.find("->")].split(",")
+            inputs = []
+            for s in subscripts:
+                shape = [dim_dict[c] for c in s]
+                inputs.append(np.ones(shape))
+            np_path = np.einsum_path(case, *inputs)
+
+            inputs = [mx.array(i) for i in inputs]
+            mx_path = mx.einsum_path(case, *inputs)
+            self.assertEqual(np_path[0][1:], mx_path[0])
+
+    def test_simple_einsum(self):
+        a = mx.arange(4 * 4).reshape(4, 4)
+        a_mx = mx.einsum("ii->i", a)
+        a_np = np.einsum("ii->i", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 2 * 2).reshape(2, 2, 2)
+        a_mx = mx.einsum("iii->i", a)
+        a_np = np.einsum("iii->i", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 2 * 3 * 3).reshape(2, 2, 3, 3)
+        a_mx = mx.einsum("iijj->ij", a)
+        a_np = np.einsum("iijj->ij", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 2 * 3 * 3).reshape(2, 3, 2, 3)
+        a_mx = mx.einsum("ijij->ij", a)
+        a_np = np.einsum("ijij->ij", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Test some simple reductions
+        a = mx.arange(2 * 2).reshape(2, 2)
+        a_mx = mx.einsum("ii", a)
+        a_np = np.einsum("ii", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 4).reshape(2, 4)
+        a_mx = mx.einsum("ij->", a)
+        a_np = np.einsum("ij->", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 4).reshape(2, 4)
+        a_mx = mx.einsum("ij->i", a)
+        a_np = np.einsum("ij->i", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 4).reshape(2, 4)
+        a_mx = mx.einsum("ij->j", a)
+        a_np = np.einsum("ij->j", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 2 * 2).reshape(2, 2, 2)
+        a_mx = mx.einsum("iii->", a)
+        a_np = np.einsum("iii->", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 2 * 3 * 3).reshape(2, 3, 2, 3)
+        a_mx = mx.einsum("ijij->j", a)
+        a_np = np.einsum("ijij->j", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Test some simple transposes
+        a = mx.arange(2 * 4).reshape(2, 4)
+        a_mx = mx.einsum("ij", a)
+        a_np = np.einsum("ij", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 4).reshape(2, 4)
+        a_mx = mx.einsum("ij->ji", a)
+        a_np = np.einsum("ij->ji", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.arange(2 * 3 * 4).reshape(2, 3, 4)
+        a_mx = mx.einsum("ijk->jki", a)
+        a_np = np.einsum("ijk->jki", a)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+    def test_two_input_einsum(self):
+
+        # Matmul
+        a = mx.full((2, 8), 1.0)
+        b = mx.full((8, 2), 1.0)
+        a_mx = mx.einsum("ik,kj", a, b)
+        a_np = np.einsum("ik,kj", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Matmul + transpose
+        a = mx.full((2, 8), 1.0)
+        b = mx.full((8, 3), 1.0)
+        a_mx = mx.einsum("ik,kj->ji", a, b)
+        a_np = np.einsum("ik,kj->ji", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Inner product
+        a = mx.full((4,), 1.0)
+        b = mx.full((4,), 1.0)
+        a_mx = mx.einsum("i,i", a, b)
+        a_np = np.einsum("i,i", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Outer product
+        a = mx.full((4,), 0.5)
+        b = mx.full((6,), 2.0)
+        a_mx = mx.einsum("i,j->ij", a, b)
+        a_np = np.einsum("i,j->ij", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Elementwise multiply
+        a = mx.full((2, 8), 1.0)
+        b = mx.full((2, 8), 1.0)
+        a_mx = mx.einsum("ij,ij->ij", a, b)
+        a_np = np.einsum("ij,ij->ij", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        # Medley
+        a = mx.full((2, 8, 3, 5), 1.0)
+        b = mx.full((3, 7, 5, 2), 1.0)
+        a_mx = mx.einsum("abcd,fgda->bfca", a, b)
+        a_np = np.einsum("abcd,fgda->bfca", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+    def test_sum_first(self):
+        a = mx.full((5, 8), 1.0)
+        b = mx.full((8, 2), 1.0)
+        a_mx = mx.einsum("ab,bc->c", a, b)
+        a_np = np.einsum("ab,bc->c", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+    def test_broadcasting(self):
+        a = mx.full((5, 1), 1.0)
+        b = mx.full((8, 2), 1.0)
+        a_mx = mx.einsum("ab,bc->c", a, b)
+        return
+        a_np = np.einsum("ab,bc->c", a, b)
+        self.assertTrue(np.array_equal(a_mx, a_np))
+
+        a = mx.random.uniform(shape=(5, 1, 3, 1))
+        b = mx.random.uniform(shape=(1, 7, 1, 2))
+        a_mx = mx.einsum("abcd,cdab->abcd", a, b)
+        a_np = np.einsum("abcd,cdab->abcd", a, b)
+        self.assertTrue(np.allclose(a_mx, a_np))
+
+    def test_attention(self):
+        q = mx.random.uniform(shape=(2, 3, 4, 5))
+        k = mx.random.uniform(shape=(2, 3, 4, 5))
+        v = mx.random.uniform(shape=(2, 3, 4, 5))
+
+        s = mx.einsum("itjk,iujk->ijtu", q, k)
+        out_mx = mx.einsum("ijtu,iujk->itjk", s, v)
+
+        s = np.einsum("itjk,iujk->ijtu", q, k)
+        out_np = np.einsum("ijtu,iujk->itjk", s, v)
+
+        self.assertTrue(np.allclose(out_mx, out_np))
+
+    def test_multi_input_einsum(self):
+        a = mx.ones((3, 4, 5))
+        out_mx = mx.einsum("ijk,lmk,ijf->lf", a, a, a)
+        out_np = np.einsum("ijk,lmk,ijf->lf", a, a, a)
+        self.assertTrue(np.allclose(out_mx, out_np))
+
+    def test_opt_einsum_test_cases(self):
+        # Test cases from
+        # https://github.com/dgasmith/opt_einsum/blob/c826bb7df16f470a69f7bf90598fc27586209d11/opt_einsum/tests/test_contract.py#L11
+        tests = [
+            # Test hadamard-like products
+            "a,ab,abc->abc",
+            "a,b,ab->ab",
+            # Test index-transformations
+            "ea,fb,gc,hd,abcd->efgh",
+            "ea,fb,abcd,gc,hd->efgh",
+            "abcd,ea,fb,gc,hd->efgh",
+            # Test complex contractions
+            "acdf,jbje,gihb,hfac,gfac,gifabc,hfac",
+            "cd,bdhe,aidb,hgca,gc,hgibcd,hgac",
+            "abhe,hidj,jgba,hiab,gab",
+            "bde,cdh,agdb,hica,ibd,hgicd,hiac",
+            "chd,bde,agbc,hiad,hgc,hgi,hiad",
+            "chd,bde,agbc,hiad,bdi,cgh,agdb",
+            "bdhe,acad,hiab,agac,hibd",
+            # Test collapse
+            "ab,ab,c->",
+            "ab,ab,c->c",
+            "ab,ab,cd,cd->",
+            "ab,ab,cd,cd->ac",
+            "ab,ab,cd,cd->cd",
+            "ab,ab,cd,cd,ef,ef->",
+            # Test outer prodcuts
+            "ab,cd,ef->abcdef",
+            "ab,cd,ef->acdf",
+            "ab,cd,de->abcde",
+            "ab,cd,de->be",
+            "ab,bcd,cd->abcd",
+            "ab,bcd,cd->abd",
+            # Random test cases that have previously failed
+            "eb,cb,fb->cef",
+            "dd,fb,be,cdb->cef",
+            "bca,cdb,dbf,afc->",
+            "dcc,fce,ea,dbf->ab",
+            "fdf,cdd,ccd,afe->ae",
+            "abcd,ad",
+            "ed,fcd,ff,bcf->be",
+            "baa,dcf,af,cde->be",
+            "bd,db,eac->ace",
+            "fff,fae,bef,def->abd",
+            "efc,dbc,acf,fd->abe",
+            # Inner products
+            "ab,ab",
+            "ab,ba",
+            "abc,abc",
+            "abc,bac",
+            "abc,cba",
+            # GEMM test cases
+            "ab,bc",
+            "ab,cb",
+            "ba,bc",
+            "ba,cb",
+            "abcd,cd",
+            "abcd,ab",
+            "abcd,cdef",
+            "abcd,cdef->feba",
+            "abcd,efdc",
+            # Inner then dot
+            "aab,bc->ac",
+            "ab,bcc->ac",
+            "aab,bcc->ac",
+            "baa,bcc->ac",
+            "aab,ccb->ac",
+            # Randomly build test caes
+            "aab,fa,df,ecc->bde",
+            "ecb,fef,bad,ed->ac",
+            "bcf,bbb,fbf,fc->",
+            "bb,ff,be->e",
+            "bcb,bb,fc,fff->",
+            "fbb,dfd,fc,fc->",
+            "afd,ba,cc,dc->bf",
+            "adb,bc,fa,cfc->d",
+            "bbd,bda,fc,db->acf",
+            "dba,ead,cad->bce",
+            "aef,fbc,dca->bde",
+        ]
+
+        size_dict = dict(zip("abcdefghij", [2, 3, 4, 5, 2, 3, 4, 5, 2, 3]))
+
+        def inputs_for_case(test_case):
+            inputs = test_case.split("->")[0].split(",")
+            return [
+                mx.random.uniform(shape=tuple(size_dict[c] for c in inp))
+                for inp in inputs
+            ]
+
+        for test_case in tests:
+            inputs = inputs_for_case(test_case)
+            np_out = np.einsum(test_case, *inputs)
+            mx_out = mx.einsum(test_case, *inputs)
+            self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4))
+
+    def test_ellipses(self):
+        size_dict = dict(zip("abcdefghij", [2, 3, 4, 5, 2, 3, 4, 5, 2, 3]))
+
+        def inputs_for_case(test_case):
+            inputs = test_case.split("->")[0].split(",")
+            return [
+                mx.random.uniform(shape=tuple(size_dict[c] for c in inp))
+                for inp in inputs
+            ]
+
+        tests = [
+            ("abc->ab", "...c->..."),
+            ("abcd->ad", "a...d->..."),
+            ("abij,abgj->abig", "...ij,...gj->...ig"),
+            ("abij,abgj->abig", "...ij,...gj->..."),
+            ("abhh->abh", "...hh->...h"),
+            ("abhh->abh", "...hh->...h"),
+            ("bch,abcj->abchj", "...h,...j->...hj"),
+            ("bc,cd->bd", "...c,cd"),
+            ("abc,acd->bd", "...bc,...cd"),
+            ("abcd,c->abd", "...cd,c"),
+            ("abcd,c->abd", "...cd,c..."),
+            ("abcd,c->abd", "...cd,c...->d..."),
+            ("abc,b->abc", "ab...,b...->ab..."),
+            ("abc,b->abc", "ab...,...b->ab..."),
+            ("abc,b->abc", "ab...,b->ab..."),
+            ("ab,bc->ac", "ab...,b...->a..."),
+            ("ab,bc->ac", "ab...,...bc->a...c"),
+            ("ab,bc->ac", "ab,b...->a..."),
+            ("abcdef,defg->abcg", "...def,defg->...g"),
+        ]
+        for test_case in tests:
+            inputs = inputs_for_case(test_case[0])
+            np_out = np.einsum(test_case[1], *inputs)
+            mx_out = mx.einsum(test_case[1], *inputs)
+            self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4))
+
+        error_tests = [
+            ("abc,abc->ab", "a...b...c,a...b...c->abc"),
+        ]
+        for test_case in error_tests:
+            inputs = inputs_for_case(test_case[0])
+            with self.assertRaises(ValueError):
+                mx.einsum(test_case[1], *inputs)
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner()

data/mlx/python/tests/test_eval.py

@@ -0,0 +1,200 @@
+# Copyright © 2023 Apple Inc.
+
+import unittest
+from functools import partial
+
+import mlx.core as mx
+import mlx_tests
+
+
+class TestEval(mlx_tests.MLXTestCase):
+    def test_eval(self):
+        arrs = [mx.ones((2, 2)) for _ in range(4)]
+        mx.eval(*arrs)
+        for x in arrs:
+            self.assertEqual(x.tolist(), [[1, 1], [1, 1]])
+
+    def test_retain_graph(self):
+        def fun(x):
+            y = 3 * x
+            mx.eval(y)
+            return 2 * y
+
+        dfun_dx = mx.grad(fun)
+        y = dfun_dx(mx.array(1.0))
+        self.assertEqual(y.item(), 6.0)
+
+    def test_eval_mixed(self):
+        x = mx.array(1) + 1 + 1
+        y = 0
+        z = "hello"
+        state = [x, y, z]
+        mx.eval(state)
+        self.assertEqual(x.item(), 3)
+
+    def test_async_eval(self):
+        x = mx.array(1) + mx.array(1) + mx.array(1)
+        mx.async_eval(x)
+        self.assertEqual(x.item(), 3)
+
+        # It should be safe to call eval on the array which has been async
+        # eval'ed
+        x = mx.array(1) + mx.array(1) + mx.array(1)
+        self.assertEqual(x.item(), 3)
+
+        x = mx.array([1, 2, 3])
+        y = 2 * x
+        mx.async_eval(y)
+        z = 2 * y
+        mx.async_eval(z)
+        self.assertTrue(mx.array_equal(y, mx.array([2, 4, 6])))
+        self.assertTrue(mx.array_equal(z, mx.array([4, 8, 12])))
+
+    def test_async_eval_twice(self):
+        for _ in range(1000):
+            x = mx.array(1) + mx.array(1) + mx.array(1)
+            mx.async_eval(x)
+            y = x + 1
+            mx.async_eval(y)
+            self.assertEqual(x.item(), 3)
+            self.assertEqual(y.item(), 4)
+
+    def test_async_eval_in_trace(self):
+        def fun(x):
+            y = x + 1.0
+            mx.async_eval(y)
+            return mx.exp(y)
+
+        # Raises
+        with self.assertRaises(ValueError):
+            mx.grad(fun)(mx.array(1.0))
+
+        # Also raises
+        with self.assertRaises(ValueError):
+            mx.vmap(fun)(mx.ones((2, 2)))
+
+    def test_async_eval_into_eval(self):
+        x = mx.array(1)
+        y = x + 1
+        mx.async_eval(y)
+        a = y - 10
+        b = mx.abs(a)
+        self.assertEqual(b.item(), 8)
+
+    def test_async_eval_into_eval_diff_stream(self):
+        s = mx.new_stream(mx.cpu)
+        x = mx.array(0)
+        y = x - 5
+        mx.async_eval(y)
+        z = mx.abs(y, stream=s)
+        self.assertEqual(z.item(), 5)
+
+    def test_eval_slow_fast_multi_stream(self):
+        x = mx.ones((8000,))
+        y = mx.abs(mx.array(-1.0))
+        for _ in range(20):
+            x = x + mx.array(1.0)
+        z = mx.add(x, y, stream=mx.cpu)
+        self.assertTrue(mx.allclose(z, mx.full((8000,), 22.0)))
+
+        # Switch eval order
+        x = mx.ones((8000,))
+        y = mx.abs(mx.array(-1.0))
+        for _ in range(20):
+            x = x + mx.array(1.0)
+        z = mx.add(y, x, stream=mx.cpu)
+        self.assertTrue(mx.allclose(z, mx.full((8000,), 22.0)))
+
+    def test_multi_output_eval_during_transform(self):
+        x = mx.random.uniform(shape=(1024,))
+        y = mx.ones((1024,))
+        mx.eval(x, y)
+
+        def fn(x):
+            a, b = mx.divmod(x, x)
+            mx.eval(a)
+            return a
+
+        out = mx.vjp(fn, (x,), (y,))
+        out = mx.vjp(fn, (x,), (y,))
+        peak_mem = mx.get_peak_memory()
+        out = mx.vjp(fn, (x,), (y,))
+        self.assertEqual(peak_mem, mx.get_peak_memory())
+
+    def test_async_eval_with_multiple_streams(self):
+        x = mx.array([1.0])
+        y = mx.array([1.0])
+        a = mx.array([1.0])
+        b = mx.array([1.0])
+
+        d = mx.default_device()
+        s2 = mx.new_stream(d)
+
+        for _ in range(50):
+            for _ in range(20):
+                x = x + y
+            mx.async_eval(x)
+            mx.eval(a + b)
+
+    def test_donation_for_noops(self):
+        def fun(x):
+            s = x.shape
+            for _ in range(10):
+                x = mx.abs(x)
+                x = mx.reshape(x, (-1,))
+                x = x.T.T
+                x = mx.stop_gradient(x)
+                x = mx.abs(x)
+            return x
+
+        x = mx.zeros((4096, 4096))
+        mx.eval(x)
+        pre = mx.get_peak_memory()
+        out = fun(x)
+        del x
+        mx.eval(out)
+        post = mx.get_peak_memory()
+        self.assertEqual(pre, post)
+
+        def fun(x):
+            for _ in range(10):
+                x = mx.abs(x)
+                x = x[:-1]
+                x = mx.abs(x)
+            return x
+
+        x = mx.zeros((4096 * 4096,))
+        mx.eval(x)
+        pre = mx.get_peak_memory()
+        out = fun(x)
+        del x
+        mx.eval(out)
+        post = mx.get_peak_memory()
+        self.assertEqual(pre, post)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available")
+    def test_multistream_deadlock(self):
+        s1 = mx.default_stream(mx.gpu)
+        s2 = mx.new_stream(mx.gpu)
+
+        x = mx.array(1.0)
+        x = mx.abs(x, stream=s1)
+        for _ in range(1000):
+            x = mx.abs(x, stream=s2)
+        mx.eval(x)
+
+        s1 = mx.default_stream(mx.gpu)
+        s2 = mx.new_stream(mx.gpu)
+        old_limit = mx.set_memory_limit(1000)
+
+        x = mx.ones((512, 512), stream=s2)
+        for _ in range(80):
+            x = mx.abs(x, stream=s1)
+        y = mx.abs(x, stream=s2)
+        z = mx.abs(y, stream=s2)
+        mx.eval(z)
+        mx.set_memory_limit(old_limit)
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner()