Reset peak memory (#1074)

* reset peak memory

* fix linux

* nits in docs

python/src/metal.cpp

@@ -35,8 +35,14 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Get the peak amount of used memory in bytes.
 
-      The maximum memory used is recorded from the beginning of the program
-      execution.
+      The maximum memory used recorded from the beginning of the program
+      execution or since the last call to :func:`reset_peak_memory`.
+      )pbdoc");
+  metal.def(
+      "reset_peak_memory",
+      &metal::reset_peak_memory,
+      R"pbdoc(
+      Reset the peak memory to zero.
       )pbdoc");
   metal.def(
       "get_cache_memory",

python/tests/test_metal.py

@@ -45,6 +45,9 @@ class TestMetal(mlx_tests.MLXTestCase):
         mx.metal.clear_cache()
         self.assertEqual(mx.metal.get_cache_memory(), 0)
 
+        mx.metal.reset_peak_memory()
+        self.assertEqual(mx.metal.get_peak_memory(), 0)
+
 
 if __name__ == "__main__":
     unittest.main()