rebase

docs/build/html/_sources/dev/extensions.rst

@@ -247,9 +247,7 @@ point-wise. This is captured in the templated function :meth:`axpby_impl`.
       float alpha_,
       float beta_,
       mx::Stream stream) {
-    // Allocate the output with `malloc_or_wait` which synchronously allocates
-    // memory, potentially waiting if the system is under memory pressure
-    out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
+    out.set_data(mx::allocator::malloc(out.nbytes()));
 
     // Get the CPU command encoder and register input and output arrays
     auto& encoder = mx::cpu::get_command_encoder(stream);
@@ -393,7 +391,7 @@ below.
         auto& d = metal::device(s.device);
 
         // Allocate output memory
-        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+        out.set_data(allocator::malloc(out.nbytes()));
 
         // Resolve name of kernel
         std::ostringstream kname;

docs/build/html/_sources/index.rst

@@ -70,6 +70,7 @@ are the CPU and GPU.
    python/fft
    python/linalg
    python/metal
+   python/memory_management
    python/nn
    python/optimizers
    python/distributed

docs/build/html/_sources/python/_autosummary/mlx.core.clear_cache.rst

@@ -0,0 +1,6 @@
+mlx.core.clear\_cache
+=====================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: clear_cache

docs/build/html/_sources/python/_autosummary/mlx.core.get_active_memory.rst

@@ -0,0 +1,6 @@
+mlx.core.get\_active\_memory
+============================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: get_active_memory

docs/build/html/_sources/python/_autosummary/mlx.core.get_cache_memory.rst

@@ -0,0 +1,6 @@
+mlx.core.get\_cache\_memory
+===========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: get_cache_memory

docs/build/html/_sources/python/_autosummary/mlx.core.get_peak_memory.rst

@@ -0,0 +1,6 @@
+mlx.core.get\_peak\_memory
+==========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: get_peak_memory

docs/build/html/_sources/python/_autosummary/mlx.core.metal.clear_cache.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.clear\_cache
-===========================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: clear_cache

docs/build/html/_sources/python/_autosummary/mlx.core.metal.get_active_memory.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.get\_active\_memory
-==================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: get_active_memory

docs/build/html/_sources/python/_autosummary/mlx.core.metal.get_cache_memory.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.get\_cache\_memory
-=================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: get_cache_memory

docs/build/html/_sources/python/_autosummary/mlx.core.metal.get_peak_memory.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.get\_peak\_memory
-================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: get_peak_memory

docs/build/html/_sources/python/_autosummary/mlx.core.metal.reset_peak_memory.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.reset\_peak\_memory
-==================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: reset_peak_memory

docs/build/html/_sources/python/_autosummary/mlx.core.metal.set_cache_limit.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.set\_cache\_limit
-================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: set_cache_limit

docs/build/html/_sources/python/_autosummary/mlx.core.metal.set_memory_limit.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.set\_memory\_limit
-=================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: set_memory_limit

docs/build/html/_sources/python/_autosummary/mlx.core.metal.set_wired_limit.rst

@@ -1,6 +0,0 @@
-mlx.core.metal.set\_wired\_limit
-================================
-
-.. currentmodule:: mlx.core.metal
-
-.. autofunction:: set_wired_limit

docs/build/html/_sources/python/_autosummary/mlx.core.reset_peak_memory.rst

@@ -0,0 +1,6 @@
+mlx.core.reset\_peak\_memory
+============================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: reset_peak_memory

docs/build/html/_sources/python/_autosummary/mlx.core.set_cache_limit.rst

@@ -0,0 +1,6 @@
+mlx.core.set\_cache\_limit
+==========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: set_cache_limit

docs/build/html/_sources/python/_autosummary/mlx.core.set_memory_limit.rst

@@ -0,0 +1,6 @@
+mlx.core.set\_memory\_limit
+===========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: set_memory_limit

docs/build/html/_sources/python/_autosummary/mlx.core.set_wired_limit.rst

@@ -0,0 +1,6 @@
+mlx.core.set\_wired\_limit
+==========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: set_wired_limit

docs/build/html/_sources/python/memory_management.rst

@@ -0,0 +1,16 @@
+Memory Management
+=================
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  get_active_memory
+  get_peak_memory
+  reset_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit
+  set_wired_limit
+  clear_cache

docs/build/html/_sources/python/metal.rst

@@ -8,13 +8,5 @@ Metal
 
   is_available
   device_info
-  get_active_memory
-  get_peak_memory
-  reset_peak_memory
-  get_cache_memory
-  set_memory_limit
-  set_cache_limit
-  set_wired_limit
-  clear_cache
   start_capture
   stop_capture