| | |
| |
|
| | import unittest |
| |
|
| | import mlx.core as mx |
| | import mlx_tests |
| |
|
| |
|
| | class TestMemory(mlx_tests.MLXTestCase): |
| | def test_memory_info(self): |
| | old_limit = mx.set_cache_limit(0) |
| |
|
| | a = mx.zeros((4096,)) |
| | mx.eval(a) |
| | del a |
| | self.assertEqual(mx.get_cache_memory(), 0) |
| | self.assertEqual(mx.set_cache_limit(old_limit), 0) |
| | self.assertEqual(mx.set_cache_limit(old_limit), old_limit) |
| |
|
| | old_limit = mx.set_memory_limit(10) |
| | self.assertTrue(mx.set_memory_limit(old_limit), 10) |
| | self.assertTrue(mx.set_memory_limit(old_limit), old_limit) |
| |
|
| | |
| | a = mx.zeros((4096,)) |
| | mx.eval(a) |
| | mx.synchronize() |
| | active_mem = mx.get_active_memory() |
| | self.assertTrue(active_mem >= 4096 * 4) |
| |
|
| | b = mx.zeros((4096,)) |
| | mx.eval(b) |
| | del b |
| | mx.synchronize() |
| |
|
| | new_active_mem = mx.get_active_memory() |
| | self.assertEqual(new_active_mem, active_mem) |
| | peak_mem = mx.get_peak_memory() |
| | self.assertTrue(peak_mem >= 4096 * 8) |
| |
|
| | if mx.metal.is_available(): |
| | cache_mem = mx.get_cache_memory() |
| | self.assertTrue(cache_mem >= 4096 * 4) |
| |
|
| | mx.clear_cache() |
| | self.assertEqual(mx.get_cache_memory(), 0) |
| |
|
| | mx.reset_peak_memory() |
| | self.assertEqual(mx.get_peak_memory(), 0) |
| |
|
| | @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") |
| | def test_wired_memory(self): |
| | old_limit = mx.set_wired_limit(1000) |
| | old_limit = mx.set_wired_limit(0) |
| | self.assertEqual(old_limit, 1000) |
| |
|
| | max_size = mx.metal.device_info()["max_recommended_working_set_size"] |
| | with self.assertRaises(ValueError): |
| | mx.set_wired_limit(max_size + 10) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | mlx_tests.MLXTestRunner() |
| |
|
