| 96 | static PyTypeObject* gc_func_tp = nullptr; |
| 97 | |
| 98 | nb::callable mlx_func( |
| 99 | nb::object func, |
| 100 | const nb::callable& orig_func, |
| 101 | std::vector<PyObject*> deps) { |
| 102 | gc_func* r = (gc_func*)PyType_GenericAlloc(gc_func_tp, 0); |
| 103 | r->func = func.inc_ref().ptr(); |
| 104 | r->orig_func = orig_func.ptr(); |
| 105 | deps.push_back(r->orig_func); |
| 106 | r->deps = std::move(deps); |
| 107 | r->vectorcall = gc_func_vectorcall; |
| 108 | return nb::steal<nb::callable>((PyObject*)r); |
| 109 | } |
| 110 | |
| 111 | void init_mlx_func(nb::module_& m) { |
| 112 | gc_func_tp = (PyTypeObject*)PyType_FromSpec(&gc_func_spec); |