| 987 | }; |
| 988 | |
| 989 | int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) { |
| 990 | Py_VISIT(Py_TYPE(self)); |
| 991 | if (!nb::inst_ready(self)) { |
| 992 | return 0; |
| 993 | } |
| 994 | |
| 995 | auto* p = nb::inst_ptr<PyCustomFunction>(self); |
| 996 | nb::handle v = nb::find(p->fun_); |
| 997 | Py_VISIT(v.ptr()); |
| 998 | if (p->vjp_fun_.has_value()) { |
| 999 | nb::handle v = nb::find(*(p->vjp_fun_)); |
| 1000 | Py_VISIT(v.ptr()); |
| 1001 | } |
| 1002 | if (p->jvp_fun_.has_value()) { |
| 1003 | nb::handle v = nb::find(*(p->jvp_fun_)); |
| 1004 | Py_VISIT(v.ptr()); |
| 1005 | } |
| 1006 | if (p->vmap_fun_.has_value()) { |
| 1007 | nb::handle v = nb::find(*(p->vmap_fun_)); |
| 1008 | Py_VISIT(v.ptr()); |
| 1009 | } |
| 1010 | return 0; |
| 1011 | } |
| 1012 | int py_custom_function_tp_clear(PyObject* self) { |
| 1013 | auto* p = nb::inst_ptr<PyCustomFunction>(self); |
| 1014 | p->reset(); |