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! 👍