zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 15:04:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Multi output primitives (#330)

Browse Source 
* Multi-output primitives

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Awni Hannun
2024-01-08 16:39:08 -08:00
committed by GitHub
parent f45f70f133
commit f099ebe535
26 changed files with 2313 additions and 1039 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/tests/test_graph.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright © 2023 Apple Inc.
import io
import unittest
import mlx.core as mx
import mlx_tests
class TestGraph(mlx_tests.MLXTestCase):
def test_to_dot(self):
# Simply test that a few cases run.
# Nothing too specific about the graph format
# for now to keep it flexible
a = mx.array(1.0)
f = io.StringIO()
mx.export_to_dot(f, a)
f.seek(0)
self.assertTrue(len(f.read()) > 0)
b = mx.array(2.0)
c = a + b
f = io.StringIO()
mx.export_to_dot(f, c)
f.seek(0)
self.assertTrue(len(f.read()) > 0)
# Multi output case
c = mx.divmod(a, b)
f = io.StringIO()
mx.export_to_dot(f, *c)
f.seek(0)
self.assertTrue(len(f.read()) > 0)
if __name__ == "__main__":
unittest.main()

24
python/tests/test_ops.py
View File

@@ -1314,7 +1314,7 @@ class TestOps(mlx_tests.MLXTestCase):
for axis in (None, 0, 1, 2):
c_npy = npop(a_npy, axis=axis)
c_mlx = mxop(a_mlx, axis=axis)
self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-4, atol=1e-4))
self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3))
for op in ["cumsum", "cumprod", "cummax", "cummin"]:
c1 = mxop(a_mlx, axis=2)
@@ -1597,6 +1597,28 @@ class TestOps(mlx_tests.MLXTestCase):
np.outer,
)
def test_divmod(self):
# A few sizes for the inputs with and without broadcasting
sizes = [
((1,), (1,)),
((1,), (10,)),
((10,), (1,)),
((3,), (3,)),
((2, 2, 2), (1, 2, 1)),
((2, 1, 2), (1, 2, 1)),
((2, 2, 2, 2), (2, 2, 2, 2)),
]
types = [np.uint16, np.uint32, np.int32, np.float16, np.float32]
for s1, s2 in sizes:
for t in types:
a_np = np.random.uniform(1, 100, size=s1).astype(t)
b_np = np.random.uniform(1, 100, size=s2).astype(t)
np_out = np.divmod(a_np, b_np)
mx_out = mx.divmod(mx.array(a_np), mx.array(b_np))
self.assertTrue(
np.allclose(np_out[0], mx_out[0]), msg=f"Shapes {s1} {s2}, Type {t}"
)
if __name__ == "__main__":
unittest.main()