Optional base class or class decorator for PyTree objects

Instead of requiring users to explicitly call register_pytree_node(), we could supply a base class that lets them instead define a pair of special methods for flattening/unflattening.

This looks slightly cleaner than calling functions on classes, but still avoids the evils of implementation inheritance.

Here’s a working example, adapted from the pytree docs:

import jax
from jax import tree_util

# baseclass that should go in tree_util.py

def _generic_flatten(tree):
  aux, children = tree.tree_flatten()
  return children, (type(tree), aux)

def _generic_unflatten(type_and_aux_data, children):
  cls, aux_data = type_and_aux_data
  return cls.tree_unflatten(aux_data, children)

# https://stackoverflow.com/questions/18126552/how-to-run-code-when-a-class-is-subclassed
class _AutoRegister(type):
  def __init__(cls, *args, **kwargs):
    tree_util.register_pytree_node(cls, _generic_flatten, _generic_unflatten)
    super().__init__(*args, **kwargs)

class PyTree(metaclass=_AutoRegister):
  def tree_flatten(self):
    """Returns aux_data and children."""
    raise NotImplemented

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    """Returns a new PyTree."""
    raise NotImplemented

# example user code

class Special(PyTree):
  def __init__(self, x, y):
    self.x = x
    self.y = y

  def __repr__(self):
    return "Special(x={}, y={})".format(self.x, self.y)

  def tree_flatten(self):
    return (None, (self.x, self.y))

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

# example usage

def show_example(structured):
  flat, tree = jax.tree_flatten(structured)
  unflattened = jax.tree_unflatten(tree, flat)
  print("structured={}\n  flat={}\n  tree={}\n  unflattened={}".format(
      structured, flat, tree, unflattened))

show_example(Special(1, 2))
# outputs:
# structured=Special(x=1, y=2)
#   flat=[1, 2]
#   tree=PyTreeDef(<class '__main__.Special'>[(<class '__main__.Special'>, None)], [*,*])
#   unflattened=Special(x=1, y=2)

[shoyer]

I like register_pytree_node_class. If it returns cls, we could use it as a decorator, too, e.g.,

@register_pytree_node_class
class Special:
  ...

It also avoids the API expansion of adding a base class that can be type checked.

[mattjj]

I think we should keep register_pytree_node_class consistent with register_pytree_node, since the former is just a tiny wrapper around the latter. If someone changes the order of register_pytree_node then both will be fixed automatically.

I don’t want to keep this issue open, so I’ll make the PR right now. I won’t update the docs, though.