| 8 | using namespace nb::literals; |
| 9 | |
| 10 | void init_memory(nb::module_& m) { |
| 11 | m.def( |
| 12 | "get_active_memory", |
| 13 | &mx::get_active_memory, |
| 14 | R"pbdoc( |
| 15 | Get the actively used memory in bytes. |
| 16 | |
| 17 | Note, this will not always match memory use reported by the system because |
| 18 | it does not include cached memory buffers. |
| 19 | )pbdoc"); |
| 20 | m.def( |
| 21 | "get_peak_memory", |
| 22 | &mx::get_peak_memory, |
| 23 | R"pbdoc( |
| 24 | Get the peak amount of used memory in bytes. |
| 25 | |
| 26 | The maximum memory used recorded from the beginning of the program |
| 27 | execution or since the last call to :func:`reset_peak_memory`. |
| 28 | )pbdoc"); |
| 29 | m.def( |
| 30 | "reset_peak_memory", |
| 31 | &mx::reset_peak_memory, |
| 32 | R"pbdoc( |
| 33 | Reset the peak memory to zero. |
| 34 | )pbdoc"); |
| 35 | m.def( |
| 36 | "get_cache_memory", |
| 37 | &mx::get_cache_memory, |
| 38 | R"pbdoc( |
| 39 | Get the cache size in bytes. |
| 40 | |
| 41 | The cache includes memory not currently used that has not been returned |
| 42 | to the system allocator. |
| 43 | )pbdoc"); |
| 44 | m.def( |
| 45 | "set_memory_limit", |
| 46 | &mx::set_memory_limit, |
| 47 | "limit"_a, |
| 48 | R"pbdoc( |
| 49 | Set the memory limit. |
| 50 | |
| 51 | The memory limit is a guideline for the maximum amount of memory to use |
| 52 | during graph evaluation. If the memory limit is exceeded and there is no |
| 53 | more RAM (including swap when available) allocations will result in an |
| 54 | exception. |
| 55 | |
| 56 | When metal is available the memory limit defaults to 1.5 times the |
| 57 | maximum recommended working set size reported by the device. |
| 58 | |
| 59 | Args: |
| 60 | limit (int): Memory limit in bytes. |
| 61 | |
| 62 | Returns: |
| 63 | int: The previous memory limit in bytes. |
| 64 | )pbdoc"); |
| 65 | m.def( |
| 66 | "set_cache_limit", |
| 67 | &mx::set_cache_limit, |