exporting a function that takes custom pytree as an argument

[frskplis]

Is it possible to serialize function that takes custom pytree as an argument?

This won't work:

from jax.tree_util import register_pytree_node_class
from jax import export

@register_pytree_node_class
class Model:
  def __init__(self, x):
    self.x = x

  def tree_flatten(self):
      children = (self.x, )
      return (children, None)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
      return cls(*children)

model = Model(1)

@jax.jit
def func(model):
  return model.x

exported: export.Exported = export.export(func)(model)
blob: bytearray = exported.serialize()

This produces:

NotImplementedError Traceback (most recent call last)
in <cell line: 1>()
----> 1 blob: bytearray = exported.serialize()
9 frames
/usr/local/lib/python3.10/dist-packages/jax/_src/export/serialization.py in _serialize_pytreedef(builder, p)
263 )
264 else:
--> 265 raise NotImplementedError(f"serializing PyTreeDef {node_data}")
266
267 ser_flatbuf.PyTreeDefStart(builder)
NotImplementedError: serializing PyTreeDef (<class 'main.Model'>, None)

[jakevdp]

You can't do this directly: exported code needs to be loadable in a different python environment, and there's no good way to serialize your class definition in a way that can be loaded in another environment (short of using something like pickle, which has other well-documented issues).

A workaround would be to define your function in terms of the flattened model buffers, and then reconstruct it inside your function:

model = Model(1)
model_flat, model_treedef = jax.tree.flatten(model)

@jax.jit
def func2(flattened_model):
  model = jax.tree.unflatten(model_treedef, flattened_model)
  return model.x

exported: export.Exported = export.export(func2)(model_flat)
blob: bytearray = exported.serialize()

[frskplis]

Thanks