zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix py gc edge case (#2079)

Browse Source
This commit is contained in:
Awni Hannun 2025-04-18 12:46:53 -07:00 committed by GitHub
parent b529515eb1
commit 55935ccae7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 10 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python/src/export.cpp
View File

@ -98,9 +98,12 @@ int py_function_exporter_tp_traverse(
PyObject* self, PyObject* self,
visitproc visit, visitproc visit,
void* arg) { void* arg) {
Py_VISIT(Py_TYPE(self));
if (!nb::inst_ready(self)) {
return 0;
}
auto* p = nb::inst_ptr<PyFunctionExporter>(self); auto* p = nb::inst_ptr<PyFunctionExporter>(self);
Py_VISIT(p->dep_.ptr()); Py_VISIT(p->dep_.ptr());
Py_VISIT(Py_TYPE(self));
return 0; return 0;
} }

2
python/src/mlx_func.cpp
View File

@ -16,12 +16,12 @@ struct gc_func {
}; };
int gc_func_tp_traverse(PyObject* self, visitproc visit, void* arg) { int gc_func_tp_traverse(PyObject* self, visitproc visit, void* arg) {
Py_VISIT(Py_TYPE(self));
gc_func* w = (gc_func*)self; gc_func* w = (gc_func*)self;
Py_VISIT(w->func); Py_VISIT(w->func);
for (auto d : w->deps) { for (auto d : w->deps) {
Py_VISIT(d); Py_VISIT(d);
} }
Py_VISIT(Py_TYPE(self));
return 0; return 0;
}; };

6
python/src/transforms.cpp
View File

@ -960,6 +960,11 @@ class PyCustomFunction {
}; };
int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) { int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) {
Py_VISIT(Py_TYPE(self));
if (!nb::inst_ready(self)) {
return 0;
}
auto* p = nb::inst_ptr<PyCustomFunction>(self); auto* p = nb::inst_ptr<PyCustomFunction>(self);
nb::handle v = nb::find(p->fun_); nb::handle v = nb::find(p->fun_);
Py_VISIT(v.ptr()); Py_VISIT(v.ptr());
@ -975,7 +980,6 @@ int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) {
nb::handle v = nb::find(*(p->vmap_fun_)); nb::handle v = nb::find(*(p->vmap_fun_));
Py_VISIT(v.ptr()); Py_VISIT(v.ptr());
} }
Py_VISIT(Py_TYPE(self));
return 0; return 0;
} }
int py_custom_function_tp_clear(PyObject* self) { int py_custom_function_tp_clear(PyObject* self) {