Repository: google/jax

Issue:
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](https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html):
```python
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)
```


Discussion:
Interesting experiment!

I disagree that calling functions on classes is bad, especially since the action being undertaken is to register a dispatch handler for a type.

I don't think we want to grow the pytree API.

Here's a simple convenience wrapper one-liner that avoids the metaclass stuff or adding the `PyTree` class to the pytree API:

```python
import operator
from jax import tree_util

def register_pytree_node_class(cls):
  flatten = operator.methodcaller('tree_flatten')
  return tree_util.register_pytree_node(cls, flatten, cls.tree_unflatten)
```

```python
# example user code

class Special:
  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 ((self.x, self.y), None)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*children)

register_pytree_node_class(Special)
```

The example usage behavior is the same. Any user who doesn't like to call registration functions on types can of course write their own subclassing magic wrapper on top!

I'm going to close this issue because it's not something we want to add to the pytree API right now.
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.
Whoa, decorator is so clever. Want to make it a PR?

Sorry if I seem like a jerk for closing the issue right after you brought it up. I'm trying to get our issues under control, since I've neglected them for too long.

Is that decorator version of `register_pytree_node_class` useful? If so, either you or I could make a PR.
No worries on the fast close, I appriecate your diligence here!

Lets reopen to consider the decorator proposal, which is definitely a better design. I should have thought of this earlier, since it's what [`flax.struct` uses](https://flax.readthedocs.io/en/latest/flax.struct.html). I think some users would find this to be a nicer model, since the class provides a nice logical grouping for these methods.

One concern is whether we should adopt the existing pytree interface, or try to fix the inconsistent order of `aux_data`/`children` between flatten/unflatten first.
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.

Resolving patch:
diff --git a/jax/tree_util.py b/jax/tree_util.py
--- a/jax/tree_util.py
+++ b/jax/tree_util.py
@@ -38,6 +38,7 @@
 
 import functools
 import collections
+import operator as op
 
 from .lib import pytree
 
@@ -105,6 +106,26 @@ def register_pytree_node(nodetype, flatten_func, unflatten_func):
   pytree.register_node(nodetype, flatten_func, unflatten_func)
   _registry[nodetype] = _RegistryEntry(flatten_func, unflatten_func)
 
+def register_pytree_node_class(cls):
+  """Extends the set of types that are considered internal nodes in pytrees.
+
+  This function is a thin wrapper around ``register_pytree_node``, and provides
+  a class-oriented interface:
+
+    @register_pytree_node_class
+    class Special:
+      def __init__(self, x, y):
+        self.x = x
+        self.y = y
+      def tree_flatten(self):
+        return ((self.x, self.y), None)
+      @classmethod
+      def tree_unflatten(cls, aux_data, children):
+        return cls(*children)
+  """
+  register_pytree_node(cls, op.methodcaller('tree_flatten'), cls.tree_unflatten)
+  return cls
+
 def tree_map(f, tree):
   """Maps a function over a pytree to produce a new pytree.
 


Test patch:
diff --git a/tests/tree_util_tests.py b/tests/tree_util_tests.py
--- a/tests/tree_util_tests.py
+++ b/tests/tree_util_tests.py
@@ -47,10 +47,28 @@ def __hash__(self):
   def __repr__(self):
     return "AnObject({},{},{})".format(self.x, self.y, self.z)
 
-
 tree_util.register_pytree_node(AnObject, lambda o: ((o.x, o.y), o.z),
                                lambda z, xy: AnObject(xy[0], xy[1], z))
 
+@tree_util.register_pytree_node_class
+class Special:
+  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 ((self.x, self.y), None)
+
+  @classmethod
+  def tree_unflatten(cls, aux_data, children):
+    return cls(*children)
+
+  def __eq__(self, other):
+    return type(self) is type(other) and (self.x, self.y) == (other.x, other.y)
+
 PYTREES = [
     ("foo",),
     ((),),
@@ -60,6 +78,7 @@ def __repr__(self):
     ([3],),
     ([3, ATuple(foo=(3, ATuple(foo=3, bar=None)), bar={"baz": 34})],),
     ([AnObject(3, None, [4, "foo"])],),
+    (Special(2, 3.),),
     ({"a": 1, "b": 2},),
     (collections.OrderedDict([("foo", 34), ("baz", 101), ("something", -42)]),),
     (collections.defaultdict(dict,