MCPcopy Create free account
hub / github.com/ml-explore/mlx / test_stream

Method test_stream

python/tests/test_device.py:68–94  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

66
67class TestStream(mlx_tests.MLXTestCase):
68 def test_stream(self):
69 s1 = mx.default_stream(mx.default_device())
70 self.assertEqual(s1.device, mx.default_device())
71
72 s2 = mx.new_stream(mx.default_device())
73 self.assertEqual(s2.device, mx.default_device())
74 self.assertNotEqual(s1, s2)
75
76 if mx.is_available(mx.gpu):
77 s_gpu = mx.default_stream(mx.gpu)
78 self.assertEqual(s_gpu.device, mx.gpu)
79 else:
80 with self.assertRaises(ValueError):
81 mx.default_stream(mx.gpu)
82
83 s_cpu = mx.default_stream(mx.cpu)
84 self.assertEqual(s_cpu.device, mx.cpu)
85
86 s_cpu = mx.new_stream(mx.cpu)
87 self.assertEqual(s_cpu.device, mx.cpu)
88
89 if mx.is_available(mx.gpu):
90 s_gpu = mx.new_stream(mx.gpu)
91 self.assertEqual(s_gpu.device, mx.gpu)
92 else:
93 with self.assertRaises(ValueError):
94 mx.new_stream(mx.gpu)
95
96 def test_op_on_stream(self):
97 x = mx.array(1.0)

Callers

nothing calls this directly

Calls 1

is_availableMethod · 0.45

Tested by

no test coverage detected