Expose function to clear memory cache (#1032)

* expose function to clear memory cache

* fix linux build

* fix metal tests

python/src/metal.cpp

@@ -90,6 +90,15 @@ void init_metal(nb::module_& m) {
       Returns:
         int: The previous cache limit in bytes.
       )pbdoc");
+  metal.def(
+      "clear_cache",
+      &metal::clear_cache,
+      R"pbdoc(
+      Clear the memory cache.
+
+      After calling this, :func:`get_cache_memory` should return ``0``.
+      )pbdoc");
+
   metal.def(
       "start_capture",
       &metal::start_capture,

python/tests/test_metal.py

@@ -42,6 +42,9 @@ class TestMetal(mlx_tests.MLXTestCase):
         cache_mem = mx.metal.get_cache_memory()
         self.assertTrue(cache_mem >= 4096 * 4)
 
+        mx.metal.clear_cache()
+        self.assertEqual(mx.metal.get_cache_memory(), 0)
+
 
 if __name__ == "__main__":
     unittest.main()