File size: 845 Bytes
712dbf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
# 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__":
mlx_tests.MLXTestRunner()
|