Repository: google/jax

Issue:
namedtuple support in arguments to transformed functions
It would be great if `xla.abstractify` would also accept namedtuples. Loop state's can consist of quite a lot of values and organizing them in a namedtuple rather than a tuple would make things nicer.


Discussion:
There's actually a convenient way to add support for custom container types throughout JAX, not just in loop carries but also for `grad`, `jit`, `vmap`, etc, all at once. Of course it's not documented at all... :)

You can register a custom type as a "pytree" (tree-like Python container) like this:

```python
from collections import namedtuple
from jax.tree_util import register_pytree_node
from jax import grad, jit
import jax.numpy as np

Point = namedtuple("Point", ["x", "y"])

register_pytree_node(
    Point,
    lambda xs: (tuple(xs), None),  # tell JAX how to unpack to an iterable
    lambda _, xs: Point(*xs)       # tell JAX how to pack back into a Point
)


def f(pt):
  return np.sqrt(pt.x**2 + pt.y**2)

pt = Point(1., 2.)

print f(pt)        # 2.236068
print grad(f)(pt)  # Point(x=..., y=...)

g = jit(f)
print g(pt)  # 2.236068
```

So that's an easy and general way to get your code working now. It also means you can have your namedtuple classes contain nested tuples/lists/dicts, or have them nested in other tuples/lists/dicts.

(By the way, the extra data that can be returned by the to-iterable function and consumed by the to-pytree fun is for things like dict keys. In the above example, we're just returning None when mapping to an iterable and then ignoring it when reconstructing.)

However, we should consider making JAX work with all namedtuple classes by default, without having to register them. Any thoughts on that, or objections to it?
I revised the issue title because we'd handle the issue in api.py and xla.abstractify would never need to see these types (just like it never sees tuples/lists/dicts).
Ha, that's awesome! Regarding namedtuple support: Given that namedtuple's are real subclasses of tuples, I think supporting all namedtuples out of the box would be the most intuitive solution.
+1 to having JAX work with all namedtuple classes
+1 Our existing codebase has been heavily relying on namedtuple and it would be great to support it in JAX.

Resolving patch:
diff --git a/jax/tree_util.py b/jax/tree_util.py
--- a/jax/tree_util.py
+++ b/jax/tree_util.py
@@ -56,7 +56,7 @@ def tree_map(f, tree):
     leaf given by `f(x)` where `x` is the value at the corresponding leaf in
     `tree`.
   """
-  node_type = node_types.get(type(tree))
+  node_type = _get_node_type(tree)
   if node_type:
     children, node_spec = node_type.to_iterable(tree)
     new_children = [tree_map(f, child) for child in children]
@@ -79,12 +79,12 @@ def tree_multimap(f, tree, *rest):
     leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf
     in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.
   """
-  node_type = node_types.get(type(tree))
+  node_type = _get_node_type(tree)
   if node_type:
     children, aux_data = node_type.to_iterable(tree)
     all_children = [children]
     for other_tree in rest:
-      other_node_type = node_types.get(type(other_tree))
+      other_node_type = _get_node_type(other_tree)
       if node_type != other_node_type:
         raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))
       other_children, other_aux_data = node_type.to_iterable(other_tree)
@@ -113,7 +113,7 @@ def process_pytree(process_node, tree):
 
 
 def walk_pytree(f_node, f_leaf, tree):
-  node_type = node_types.get(type(tree))
+  node_type = _get_node_type(tree)
   if node_type:
     children, node_spec = node_type.to_iterable(tree)
     proc_children, child_specs = unzip2([walk_pytree(f_node, f_leaf, child)
@@ -236,3 +236,20 @@ def register_pytree_node(py_type, to_iterable, from_iterable):
 register_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))
 register_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))
 register_pytree_node(type(None), lambda z: ((), None), lambda _, xs: None)
+
+
+# To handle namedtuples, we can't just use the standard table of node_types
+# because every namedtuple creates its own type and thus would require its own
+# entry in the table. Instead we use a heuristic check on the type itself to
+# decide whether it's a namedtuple type, and if so treat it as a pytree node.
+def _get_node_type(maybe_tree):
+  t = type(maybe_tree)
+  return node_types.get(t) or _namedtuple_node(t)
+
+def _namedtuple_node(t):
+  if t.__bases__ == (tuple,) and hasattr(t, '_fields'):
+    return NamedtupleNode
+
+NamedtupleNode = NodeType('namedtuple',
+                          lambda xs: (tuple(xs), type(xs)),
+                          lambda t, xs: t(*xs))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -16,11 +16,11 @@
 from __future__ import division
 from __future__ import print_function
 
-import six
+import collections
 
-import numpy as onp
 from absl.testing import absltest
-from jax import test_util as jtu
+import numpy as onp
+import six
 
 import jax.numpy as np
 from jax import jit, grad, device_get, device_put, jacfwd, jacrev, hessian
@@ -29,6 +29,7 @@
 from jax.interpreters.ad import defjvp, defvjp, defvjp2, defvjp_all
 from jax.interpreters.xla import DeviceArray, DeviceTuple
 from jax.abstract_arrays import concretization_err_msg
+from jax import test_util as jtu
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -568,6 +569,22 @@ def test_devicearray_repr(self):
     self.assertIsInstance(x, DeviceArray)
     repr(x)  # doesn't crash
 
+  def test_namedtuple_transparency(self):
+    # See https://github.com/google/jax/issues/446
+    Point = collections.namedtuple("Point", ["x", "y"])
+
+    def f(pt):
+      return np.sqrt(pt.x ** 2 + pt.y ** 2)
+
+    pt = Point(1., 2.)
+
+    f(pt)  # doesn't crash
+    g = api.grad(f)(pt)
+    self.assertIsInstance(g, Point)
+
+    f_jit = api.jit(f)
+    self.assertAllClose(f(pt), f_jit(pt), check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()