zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Cycle leak break (#1856)

Browse Source 
* detect and break leaks in custom function

* detect and break leaks in custom function
This commit is contained in:
Awni Hannun
2025-02-11 14:45:02 -08:00
committed by GitHub
parent 142b77751d
commit 2a45056ba8
10 changed files with 396 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/src/mlx_func.h Normal file
View File

@@ -0,0 +1,24 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <vector>
#include <nanobind/nanobind.h>
#include <nanobind/stl/function.h>
namespace nb = nanobind;
using namespace nb::literals;
nb::callable mlx_func(nb::object func, std::vector<PyObject*> deps);
template <typename F, typename... Deps>
nb::callable mlx_func(F func, Deps&&... deps) {
return mlx_func(
nb::cpp_function(std::move(func)), std::vector<PyObject*>{deps.ptr()...});
}
template <typename... Deps>
nb::callable mlx_func(nb::object func, Deps&&... deps) {
return mlx_func(std::move(func), std::vector<PyObject*>{deps.ptr()...});
}