Pickling nanobind functions (nb_func)

[awni]

Just curious on how possible to have support for this or to extend nb_func to allow it to be pickled?

This comes up in MLX since compiled functions are nb_func. And sometimes it's convenient for a Python class to hold a function as an attribute. We can work around it by wrapping stuff in user defined Python classes.. but it's inconvenient and confusing for users who do not know about this.

[wjakob]

Sorry for the late reply @awni . What would it mean to pickle a nb_func? It contains arbitrary internal state (pointers to functions), closure capture objects (which themselves can contain pointers).

While that can probably be serialized to a file, I don't see how unserializing it would produce a meaningful result.

[awni]

Apologies, it was not a very well thought out question. And it does seem like the nb_func is not the right place to solve this issue.

Let me explain the problem.. and if you have any suggestions I would love to hear them and if not no worries:

@mx.compile
def myfun(x):
   # do some stuff

class MyClass:
  def __init__(self):
    self.fn = myfun

Normally you could pickle MyClass but the output of mx.compile(myfun) is a nb_func, so it breaks that.

[wjakob]

Did you not implement your own GCable nb_func wrappers at some point? If so, you could likely install custom pickling/unpickling handers that wrap pickling of the underlying Python callable and reapply @mx.compile when unpickled.

[awni]

Did you not implement your own GCable nb_func wrappers at some point?

Indeed I did implement those. So yea the output of mx.compile is actually a wrapper of nb_func.

If so, you could likely install custom pickling/unpickling handers that wrap pickling of the underlying Python callable and reapply @mx.compile when unpickled.

Makes sense! 👍

[oremanj]

I'm carrying a local patch that lets nb_funcs (and methods) be pickled. It's not difficult if you pickle them like Python functions are pickled, i.e., by their name rather than by their contents.

diff --git a/atl/src/3ps/nanobind/src/nb_func.cpp b/atl/src/3ps/nanobind/src/nb_func.cpp
index b2294629f0adc..42be0fa791c8e 100644
--- a/atl/src/3ps/nanobind/src/nb_func.cpp
+++ b/atl/src/3ps/nanobind/src/nb_func.cpp
@@ -1066,6 +1066,10 @@ PyObject *nb_func_get_doc(PyObject *self, void *) {
     return PyUnicode_FromString(buf.get());
 }

+PyObject *nb_func_reduce(PyObject *self, PyObject *args_unused) {
+    return nb_func_get_qualname(self);
+}
+
 // PyGetSetDef entry for __module__ is ignored in Python 3.8
 PyObject *nb_func_getattro(PyObject *self, PyObject *name_) {
     const char *name = PyUnicode_AsUTF8AndSize(name_, nullptr);
@@ -1084,9 +1088,47 @@ PyObject *nb_func_getattro(PyObject *self, PyObject *name_) {
         return PyObject_GenericGetAttr(self, name_);
 }

-PyObject *nb_bound_method_getattro(PyObject *self, PyObject *name) {
-    nb_func *func = ((nb_bound_method *) self)->func;
-    return nb_func_getattro((PyObject *) func, name);
+PyObject *nb_bound_method_getattro(PyObject *self, PyObject *name_) {
+    const char *name = PyUnicode_AsUTF8AndSize(name_, nullptr);
+
+    if (!name)
+        return nullptr;
+    else if (strcmp(name, "__reduce__") == 0 ||
+             strcmp(name, "__reduce_ex__") == 0)
+        return PyObject_GenericGetAttr(self, name_);
+    else if (strcmp(name, "__self__") == 0) {
+        PyObject *mb_self = ((nb_bound_method *) self)->self;
+        Py_INCREF(mb_self);
+        return mb_self;
+    } else if (strcmp(name, "__func__") == 0) {
+        PyObject *mb_func = (PyObject *) (((nb_bound_method *) self)->func);
+        Py_INCREF(mb_func);
+        return mb_func;
+    } else {
+        nb_func *func = ((nb_bound_method *) self)->func;
+        return nb_func_getattro((PyObject *) func, name_);
+    }
+}
+
+PyObject *nb_bound_method_reduce(PyObject *self, PyObject *args_unused) {
+    nb_bound_method *mb = (nb_bound_method *) self;
+    PyObject *name = nb_func_get_name((PyObject *) mb->func);
+    if (name == nullptr)
+        return nullptr;
+    if (name == Py_None) {
+        Py_DECREF(name);
+        PyErr_Format(PyExc_TypeError,
+                     "can't pickle unnamed nb_bound_method of %R", mb->self);
+        return nullptr;
+    }
+    PyObject *builtins = PyEval_GetBuiltins(); // borrowed ref + never fails
+    PyObject *getattr = PyDict_GetItemString(builtins, "getattr");
+    if (getattr == nullptr) {
+        Py_DECREF(name);
+        PyErr_SetString(PyExc_SystemError, "can't find 'getattr' builtin");
+        return nullptr;
+    }
+    return Py_BuildValue("N(ON)", getattr, mb->self, name);
 }

 /// Excise a substring from 's'
diff --git a/atl/src/3ps/nanobind/src/nb_internals.cpp b/atl/src/3ps/nanobind/src/nb_internals.cpp
index 18fa008247f8d..99c36c2831219 100644
--- a/atl/src/3ps/nanobind/src/nb_internals.cpp
+++ b/atl/src/3ps/nanobind/src/nb_internals.cpp
@@ -79,6 +79,8 @@ extern PyObject *nb_func_getattro(PyObject *, PyObject *);
 extern PyObject *nb_func_get_doc(PyObject *, void *);
 extern PyObject *nb_bound_method_getattro(PyObject *, PyObject *);
 extern PyObject *nb_func_get_getset(PyObject *, PyObject *);
+extern PyObject *nb_func_reduce(PyObject *, PyObject *);
+extern PyObject *nb_bound_method_reduce(PyObject *, PyObject *);
 extern int nb_func_traverse(PyObject *, visitproc, void *);
 extern int nb_func_clear(PyObject *);
 extern void nb_func_dealloc(PyObject *);
@@ -110,8 +112,14 @@ static PyGetSetDef nb_func_getset[] = {
     { nullptr, nullptr, nullptr, nullptr, nullptr }
 };

+static PyMethodDef nb_func_methods[] = {
+    { "__reduce__", nb_func_reduce, METH_NOARGS, nullptr },
+    { nullptr, nullptr, 0, nullptr }
+};
+
 static PyType_Slot nb_func_slots[] = {
     { Py_tp_members, (void *) nb_func_members },
+    { Py_tp_methods, (void *) nb_func_methods },
     { Py_tp_getset, (void *) nb_func_getset },
     { Py_tp_getattro, (void *) nb_func_getattro },
     { Py_tp_traverse, (void *) nb_func_traverse },
@@ -134,6 +142,7 @@ static PyType_Spec nb_func_spec = {

 static PyType_Slot nb_method_slots[] = {
     { Py_tp_members, (void *) nb_func_members },
+    { Py_tp_methods, (void *) nb_func_methods },
     { Py_tp_getset, (void *) nb_func_getset },
     { Py_tp_getattro, (void *) nb_func_getattro },
     { Py_tp_traverse, (void *) nb_func_traverse },
@@ -161,8 +170,14 @@ static PyMemberDef nb_bound_method_members[] = {
     { nullptr, 0, 0, 0, nullptr }
 };

+static PyMethodDef nb_bound_method_methods[] = {
+    { "__reduce__", nb_bound_method_reduce, METH_NOARGS, nullptr, },
+    { nullptr, nullptr, 0, nullptr }
+};
+
 static PyType_Slot nb_bound_method_slots[] = {
     { Py_tp_members, (void *) nb_bound_method_members },
+    { Py_tp_methods, (void *) nb_bound_method_methods },
     { Py_tp_getattro, (void *) nb_bound_method_getattro },
     { Py_tp_traverse, (void *) nb_bound_method_traverse },
     { Py_tp_clear, (void *) nb_bound_method_clear },

Not sure if this is considered worthwhile upstream but I'd be happy to submit a PR if so.