We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/tree_util.py`
Content:
```
1 # Copyright 2018 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Utilities for working with tree-like container data structures.
16 
17 The code here is independent of JAX. The only dependence is on jax.util, which
18 itself has no JAX-specific code.
19 
20 This module provides a small set of utility functions for working with tree-like
21 data structures, such as nested tuples, lists, and dicts. We call these
22 structures pytrees. They are trees in that they are defined recursively (any
23 non-pytree is a pytree, i.e. a leaf, and any pytree of pytrees is a pytree) and
24 can be operated on recursively (object identity equivalence is not preserved by
25 mapping operations, and the structures cannot contain reference cycles).
26 
27 The set of Python types that are considered pytree nodes (e.g. that can be
28 mapped over, rather than treated as leaves) is extensible. There is a single
29 module-level registry of types, and class hierarchy is ignored. By registering a
30 new pytree node type, that type in effect becomes transparent to the utility
31 functions in this file.
32 """
33 
34 from __future__ import absolute_import
35 from __future__ import division
36 from __future__ import print_function
37 
38 from collections import namedtuple
39 import itertools as it
40 from six.moves import reduce
41 
42 from .util import unzip2, concatenate, partial, safe_map
43 
44 map = safe_map
45 
46 
47 def tree_map(f, tree):
48   """Map a function over a pytree to produce a new pytree.
49 
50   Args:
51     f: function to be applied at each leaf.
52     tree: a pytree to be mapped over.
53 
54   Returns:
55     A new pytree with the same structure as `tree` but with the value at each
56     leaf given by `f(x)` where `x` is the value at the corresponding leaf in
57     `tree`.
58   """
59   node_type = node_types.get(type(tree))
60   if node_type:
61     children, node_spec = node_type.to_iterable(tree)
62     new_children = [tree_map(f, child) for child in children]
63     return node_type.from_iterable(node_spec, new_children)
64   else:
65     return f(tree)
66 
67 def tree_multimap(f, tree, *rest):
68   """Map a multi-input function over pytree args to produce a new pytree.
69 
70   Args:
71     f: function that takes `1 + len(rest)` arguments, to be applied at the
72       corresponding leaves of the pytrees.
73     tree: a pytree to be mapped over, with each leaf providing the first
74       positional argument to `f`.
75     *rest: a tuple of pytrees, each with the same structure as `tree`.
76 
77   Returns:
78     A new pytree with the same structure as `tree` but with the value at each
79     leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf
80     in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.
81   """
82   node_type = node_types.get(type(tree))
83   if node_type:
84     children, aux_data = node_type.to_iterable(tree)
85     all_children = [children]
86     for other_tree in rest:
87       other_node_type = node_types.get(type(other_tree))
88       if node_type != other_node_type:
89         raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))
90       other_children, other_aux_data = node_type.to_iterable(other_tree)
91       if other_aux_data != aux_data:
92         raise TypeError('Mismatch: {} != {}'.format(other_aux_data, aux_data))
93       all_children.append(other_children)
94 
95     new_children = [tree_multimap(f, *xs) for xs in zip(*all_children)]
96     return node_type.from_iterable(aux_data, new_children)
97   else:
98     return f(tree, *rest)
99 
100 
101 def tree_reduce(f, tree):
102   flat, _ = tree_flatten(tree)
103   return reduce(f, flat)
104 
105 
106 def tree_all(tree):
107   flat, _ = tree_flatten(tree)
108   return all(flat)
109 
110 
111 def process_pytree(process_node, tree):
112   return walk_pytree(process_node, lambda x: x, tree)
113 
114 
115 def walk_pytree(f_node, f_leaf, tree):
116   node_type = node_types.get(type(tree))
117   if node_type:
118     children, node_spec = node_type.to_iterable(tree)
119     proc_children, child_specs = unzip2([walk_pytree(f_node, f_leaf, child)
120                                          for child in children])
121     tree_def = PyTreeDef(node_type, node_spec, child_specs)
122     return f_node(proc_children), tree_def
123   else:
124     return f_leaf(tree), leaf
125 
126 
127 def build_tree(treedef, xs):
128   if treedef is leaf:
129     return xs
130   else:
131     # We use 'iter' for clearer error messages
132     children = map(build_tree, iter(treedef.children), iter(xs))
133     return treedef.node_type.from_iterable(treedef.node_data, children)
134 
135 
136 tree_flatten = partial(walk_pytree, concatenate, lambda x: [x])
137 
138 def tree_unflatten(treedef, xs):
139   return _tree_unflatten(iter(xs), treedef)
140 
141 def _tree_unflatten(xs, treedef):
142   if treedef is leaf:
143     return next(xs)
144   else:
145     children = map(partial(_tree_unflatten, xs), treedef.children)
146     return treedef.node_type.from_iterable(treedef.node_data, children)
147 
148 
149 def tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):
150   flat, treedef = tree_flatten(pytree_to_transpose)
151   expected_treedef = _nested_treedef(inner_treedef, outer_treedef)
152   if treedef != expected_treedef:
153     raise TypeError("Mismatch\n{}\n != \n{}".format(treedef, expected_treedef))
154 
155   inner_size = _num_leaves(inner_treedef)
156   outer_size = _num_leaves(outer_treedef)
157   flat = iter(flat)
158   lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]
159   transposed_lol = zip(*lol)
160   subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)
161   return tree_unflatten(inner_treedef, subtrees)
162 
163 def _num_leaves(treedef):
164   return 1 if treedef is leaf else sum(map(_num_leaves, treedef.children))
165 
166 def _nested_treedef(inner, outer):
167   # just used in tree_transpose error checking
168   if outer is leaf:
169     return inner
170   else:
171     children = map(partial(_nested_treedef, inner), outer.children)
172     return PyTreeDef(outer.node_type, outer.node_data, tuple(children))
173 
174 
175 def tree_structure(tree):
176   _, spec = process_pytree(lambda _: None, tree)
177   return spec
178 
179 
180 class PyTreeDef(object):
181   def __init__(self, node_type, node_data, children):
182     self.node_type = node_type
183     self.node_data = node_data
184     self.children = children
185 
186   def __repr__(self):
187     if self.node_data is None:
188       data_repr = ""
189     else:
190       data_repr = "[{}]".format(self.node_data)
191 
192     return "PyTree({}{}, [{}])".format(self.node_type.name, data_repr,
193                                      ','.join(map(repr, self.children)))
194 
195   def __hash__(self):
196     return hash((self.node_type, self.node_data, tuple(self.children)))
197 
198   def __eq__(self, other):
199     if other is leaf:
200       return False
201     else:
202       return (self.node_type == other.node_type and
203               self.node_data == other.node_data and
204               self.children == other.children)
205 
206   def __ne__(self, other):
207     return not self == other
208 
209 
210 class PyLeaf(object):
211   def __repr__(self):
212     return '*'
213 
214 leaf = PyLeaf()
215 
216 def dict_to_iterable(xs):
217   keys = tuple(sorted(xs.keys()))
218   return tuple(map(xs.get, keys)), keys
219 
220 class NodeType(object):
221   def __init__(self, name, to_iterable, from_iterable):
222     self.name = name
223     self.to_iterable = to_iterable
224     self.from_iterable = from_iterable
225 
226   def __repr__(self):
227     return self.name
228 
229 node_types = {}
230 
231 def register_pytree_node(py_type, to_iterable, from_iterable):
232   assert py_type not in node_types
233   node_types[py_type] = NodeType(str(py_type), to_iterable, from_iterable)
234 
235 register_pytree_node(tuple, lambda xs: (xs, None), lambda _, xs: tuple(xs))
236 register_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))
237 register_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))
238 register_pytree_node(type(None), lambda z: ((), None), lambda _, xs: None)
239
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
     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 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 @@
 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))

{"golden_diff": "diff --git a/jax/tree_util.py b/jax/tree_util.py\n--- a/jax/tree_util.py\n+++ b/jax/tree_util.py\n@@ -56,7 +56,7 @@\n     leaf given by `f(x)` where `x` is the value at the corresponding leaf in\n     `tree`.\n   \"\"\"\n-  node_type = node_types.get(type(tree))\n+  node_type = _get_node_type(tree)\n   if node_type:\n     children, node_spec = node_type.to_iterable(tree)\n     new_children = [tree_map(f, child) for child in children]\n@@ -79,12 +79,12 @@\n     leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf\n     in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.\n   \"\"\"\n-  node_type = node_types.get(type(tree))\n+  node_type = _get_node_type(tree)\n   if node_type:\n     children, aux_data = node_type.to_iterable(tree)\n     all_children = [children]\n     for other_tree in rest:\n-      other_node_type = node_types.get(type(other_tree))\n+      other_node_type = _get_node_type(other_tree)\n       if node_type != other_node_type:\n         raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))\n       other_children, other_aux_data = node_type.to_iterable(other_tree)\n@@ -113,7 +113,7 @@\n \n \n def walk_pytree(f_node, f_leaf, tree):\n-  node_type = node_types.get(type(tree))\n+  node_type = _get_node_type(tree)\n   if node_type:\n     children, node_spec = node_type.to_iterable(tree)\n     proc_children, child_specs = unzip2([walk_pytree(f_node, f_leaf, child)\n@@ -236,3 +236,20 @@\n register_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))\n register_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))\n register_pytree_node(type(None), lambda z: ((), None), lambda _, xs: None)\n+\n+\n+# To handle namedtuples, we can't just use the standard table of node_types\n+# because every namedtuple creates its own type and thus would require its own\n+# entry in the table. Instead we use a heuristic check on the type itself to\n+# decide whether it's a namedtuple type, and if so treat it as a pytree node.\n+def _get_node_type(maybe_tree):\n+  t = type(maybe_tree)\n+  return node_types.get(t) or _namedtuple_node(t)\n+\n+def _namedtuple_node(t):\n+  if t.__bases__ == (tuple,) and hasattr(t, '_fields'):\n+    return NamedtupleNode\n+\n+NamedtupleNode = NodeType('namedtuple',\n+                          lambda xs: (tuple(xs), type(xs)),\n+                          lambda t, xs: t(*xs))\n", "issue": "namedtuple support in arguments to transformed functions\nIt 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.\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for working with tree-like container data structures.\n\nThe code here is independent of JAX. The only dependence is on jax.util, which\nitself has no JAX-specific code.\n\nThis module provides a small set of utility functions for working with tree-like\ndata structures, such as nested tuples, lists, and dicts. We call these\nstructures pytrees. They are trees in that they are defined recursively (any\nnon-pytree is a pytree, i.e. a leaf, and any pytree of pytrees is a pytree) and\ncan be operated on recursively (object identity equivalence is not preserved by\nmapping operations, and the structures cannot contain reference cycles).\n\nThe set of Python types that are considered pytree nodes (e.g. that can be\nmapped over, rather than treated as leaves) is extensible. There is a single\nmodule-level registry of types, and class hierarchy is ignored. By registering a\nnew pytree node type, that type in effect becomes transparent to the utility\nfunctions in this file.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom collections import namedtuple\nimport itertools as it\nfrom six.moves import reduce\n\nfrom .util import unzip2, concatenate, partial, safe_map\n\nmap = safe_map\n\n\ndef tree_map(f, tree):\n  \"\"\"Map a function over a pytree to produce a new pytree.\n\n  Args:\n    f: function to be applied at each leaf.\n    tree: a pytree to be mapped over.\n\n  Returns:\n    A new pytree with the same structure as `tree` but with the value at each\n    leaf given by `f(x)` where `x` is the value at the corresponding leaf in\n    `tree`.\n  \"\"\"\n  node_type = node_types.get(type(tree))\n  if node_type:\n    children, node_spec = node_type.to_iterable(tree)\n    new_children = [tree_map(f, child) for child in children]\n    return node_type.from_iterable(node_spec, new_children)\n  else:\n    return f(tree)\n\ndef tree_multimap(f, tree, *rest):\n  \"\"\"Map a multi-input function over pytree args to produce a new pytree.\n\n  Args:\n    f: function that takes `1 + len(rest)` arguments, to be applied at the\n      corresponding leaves of the pytrees.\n    tree: a pytree to be mapped over, with each leaf providing the first\n      positional argument to `f`.\n    *rest: a tuple of pytrees, each with the same structure as `tree`.\n\n  Returns:\n    A new pytree with the same structure as `tree` but with the value at each\n    leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf\n    in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.\n  \"\"\"\n  node_type = node_types.get(type(tree))\n  if node_type:\n    children, aux_data = node_type.to_iterable(tree)\n    all_children = [children]\n    for other_tree in rest:\n      other_node_type = node_types.get(type(other_tree))\n      if node_type != other_node_type:\n        raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))\n      other_children, other_aux_data = node_type.to_iterable(other_tree)\n      if other_aux_data != aux_data:\n        raise TypeError('Mismatch: {} != {}'.format(other_aux_data, aux_data))\n      all_children.append(other_children)\n\n    new_children = [tree_multimap(f, *xs) for xs in zip(*all_children)]\n    return node_type.from_iterable(aux_data, new_children)\n  else:\n    return f(tree, *rest)\n\n\ndef tree_reduce(f, tree):\n  flat, _ = tree_flatten(tree)\n  return reduce(f, flat)\n\n\ndef tree_all(tree):\n  flat, _ = tree_flatten(tree)\n  return all(flat)\n\n\ndef process_pytree(process_node, tree):\n  return walk_pytree(process_node, lambda x: x, tree)\n\n\ndef walk_pytree(f_node, f_leaf, tree):\n  node_type = node_types.get(type(tree))\n  if node_type:\n    children, node_spec = node_type.to_iterable(tree)\n    proc_children, child_specs = unzip2([walk_pytree(f_node, f_leaf, child)\n                                         for child in children])\n    tree_def = PyTreeDef(node_type, node_spec, child_specs)\n    return f_node(proc_children), tree_def\n  else:\n    return f_leaf(tree), leaf\n\n\ndef build_tree(treedef, xs):\n  if treedef is leaf:\n    return xs\n  else:\n    # We use 'iter' for clearer error messages\n    children = map(build_tree, iter(treedef.children), iter(xs))\n    return treedef.node_type.from_iterable(treedef.node_data, children)\n\n\ntree_flatten = partial(walk_pytree, concatenate, lambda x: [x])\n\ndef tree_unflatten(treedef, xs):\n  return _tree_unflatten(iter(xs), treedef)\n\ndef _tree_unflatten(xs, treedef):\n  if treedef is leaf:\n    return next(xs)\n  else:\n    children = map(partial(_tree_unflatten, xs), treedef.children)\n    return treedef.node_type.from_iterable(treedef.node_data, children)\n\n\ndef tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):\n  flat, treedef = tree_flatten(pytree_to_transpose)\n  expected_treedef = _nested_treedef(inner_treedef, outer_treedef)\n  if treedef != expected_treedef:\n    raise TypeError(\"Mismatch\\n{}\\n != \\n{}\".format(treedef, expected_treedef))\n\n  inner_size = _num_leaves(inner_treedef)\n  outer_size = _num_leaves(outer_treedef)\n  flat = iter(flat)\n  lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]\n  transposed_lol = zip(*lol)\n  subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)\n  return tree_unflatten(inner_treedef, subtrees)\n\ndef _num_leaves(treedef):\n  return 1 if treedef is leaf else sum(map(_num_leaves, treedef.children))\n\ndef _nested_treedef(inner, outer):\n  # just used in tree_transpose error checking\n  if outer is leaf:\n    return inner\n  else:\n    children = map(partial(_nested_treedef, inner), outer.children)\n    return PyTreeDef(outer.node_type, outer.node_data, tuple(children))\n\n\ndef tree_structure(tree):\n  _, spec = process_pytree(lambda _: None, tree)\n  return spec\n\n\nclass PyTreeDef(object):\n  def __init__(self, node_type, node_data, children):\n    self.node_type = node_type\n    self.node_data = node_data\n    self.children = children\n\n  def __repr__(self):\n    if self.node_data is None:\n      data_repr = \"\"\n    else:\n      data_repr = \"[{}]\".format(self.node_data)\n\n    return \"PyTree({}{}, [{}])\".format(self.node_type.name, data_repr,\n                                     ','.join(map(repr, self.children)))\n\n  def __hash__(self):\n    return hash((self.node_type, self.node_data, tuple(self.children)))\n\n  def __eq__(self, other):\n    if other is leaf:\n      return False\n    else:\n      return (self.node_type == other.node_type and\n              self.node_data == other.node_data and\n              self.children == other.children)\n\n  def __ne__(self, other):\n    return not self == other\n\n\nclass PyLeaf(object):\n  def __repr__(self):\n    return '*'\n\nleaf = PyLeaf()\n\ndef dict_to_iterable(xs):\n  keys = tuple(sorted(xs.keys()))\n  return tuple(map(xs.get, keys)), keys\n\nclass NodeType(object):\n  def __init__(self, name, to_iterable, from_iterable):\n    self.name = name\n    self.to_iterable = to_iterable\n    self.from_iterable = from_iterable\n\n  def __repr__(self):\n    return self.name\n\nnode_types = {}\n\ndef register_pytree_node(py_type, to_iterable, from_iterable):\n  assert py_type not in node_types\n  node_types[py_type] = NodeType(str(py_type), to_iterable, from_iterable)\n\nregister_pytree_node(tuple, lambda xs: (xs, None), lambda _, xs: tuple(xs))\nregister_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))\nregister_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))\nregister_pytree_node(type(None), lambda z: ((), None), lambda _, xs: None)\n", "path": "jax/tree_util.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for working with tree-like container data structures.\n\nThe code here is independent of JAX. The only dependence is on jax.util, which\nitself has no JAX-specific code.\n\nThis module provides a small set of utility functions for working with tree-like\ndata structures, such as nested tuples, lists, and dicts. We call these\nstructures pytrees. They are trees in that they are defined recursively (any\nnon-pytree is a pytree, i.e. a leaf, and any pytree of pytrees is a pytree) and\ncan be operated on recursively (object identity equivalence is not preserved by\nmapping operations, and the structures cannot contain reference cycles).\n\nThe set of Python types that are considered pytree nodes (e.g. that can be\nmapped over, rather than treated as leaves) is extensible. There is a single\nmodule-level registry of types, and class hierarchy is ignored. By registering a\nnew pytree node type, that type in effect becomes transparent to the utility\nfunctions in this file.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom collections import namedtuple\nimport itertools as it\nfrom six.moves import reduce\n\nfrom .util import unzip2, concatenate, partial, safe_map\n\nmap = safe_map\n\n\ndef tree_map(f, tree):\n  \"\"\"Map a function over a pytree to produce a new pytree.\n\n  Args:\n    f: function to be applied at each leaf.\n    tree: a pytree to be mapped over.\n\n  Returns:\n    A new pytree with the same structure as `tree` but with the value at each\n    leaf given by `f(x)` where `x` is the value at the corresponding leaf in\n    `tree`.\n  \"\"\"\n  node_type = _get_node_type(tree)\n  if node_type:\n    children, node_spec = node_type.to_iterable(tree)\n    new_children = [tree_map(f, child) for child in children]\n    return node_type.from_iterable(node_spec, new_children)\n  else:\n    return f(tree)\n\ndef tree_multimap(f, tree, *rest):\n  \"\"\"Map a multi-input function over pytree args to produce a new pytree.\n\n  Args:\n    f: function that takes `1 + len(rest)` arguments, to be applied at the\n      corresponding leaves of the pytrees.\n    tree: a pytree to be mapped over, with each leaf providing the first\n      positional argument to `f`.\n    *rest: a tuple of pytrees, each with the same structure as `tree`.\n\n  Returns:\n    A new pytree with the same structure as `tree` but with the value at each\n    leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf\n    in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.\n  \"\"\"\n  node_type = _get_node_type(tree)\n  if node_type:\n    children, aux_data = node_type.to_iterable(tree)\n    all_children = [children]\n    for other_tree in rest:\n      other_node_type = _get_node_type(other_tree)\n      if node_type != other_node_type:\n        raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))\n      other_children, other_aux_data = node_type.to_iterable(other_tree)\n      if other_aux_data != aux_data:\n        raise TypeError('Mismatch: {} != {}'.format(other_aux_data, aux_data))\n      all_children.append(other_children)\n\n    new_children = [tree_multimap(f, *xs) for xs in zip(*all_children)]\n    return node_type.from_iterable(aux_data, new_children)\n  else:\n    return f(tree, *rest)\n\n\ndef tree_reduce(f, tree):\n  flat, _ = tree_flatten(tree)\n  return reduce(f, flat)\n\n\ndef tree_all(tree):\n  flat, _ = tree_flatten(tree)\n  return all(flat)\n\n\ndef process_pytree(process_node, tree):\n  return walk_pytree(process_node, lambda x: x, tree)\n\n\ndef walk_pytree(f_node, f_leaf, tree):\n  node_type = _get_node_type(tree)\n  if node_type:\n    children, node_spec = node_type.to_iterable(tree)\n    proc_children, child_specs = unzip2([walk_pytree(f_node, f_leaf, child)\n                                         for child in children])\n    tree_def = PyTreeDef(node_type, node_spec, child_specs)\n    return f_node(proc_children), tree_def\n  else:\n    return f_leaf(tree), leaf\n\n\ndef build_tree(treedef, xs):\n  if treedef is leaf:\n    return xs\n  else:\n    # We use 'iter' for clearer error messages\n    children = map(build_tree, iter(treedef.children), iter(xs))\n    return treedef.node_type.from_iterable(treedef.node_data, children)\n\n\ntree_flatten = partial(walk_pytree, concatenate, lambda x: [x])\n\ndef tree_unflatten(treedef, xs):\n  return _tree_unflatten(iter(xs), treedef)\n\ndef _tree_unflatten(xs, treedef):\n  if treedef is leaf:\n    return next(xs)\n  else:\n    children = map(partial(_tree_unflatten, xs), treedef.children)\n    return treedef.node_type.from_iterable(treedef.node_data, children)\n\n\ndef tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):\n  flat, treedef = tree_flatten(pytree_to_transpose)\n  expected_treedef = _nested_treedef(inner_treedef, outer_treedef)\n  if treedef != expected_treedef:\n    raise TypeError(\"Mismatch\\n{}\\n != \\n{}\".format(treedef, expected_treedef))\n\n  inner_size = _num_leaves(inner_treedef)\n  outer_size = _num_leaves(outer_treedef)\n  flat = iter(flat)\n  lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]\n  transposed_lol = zip(*lol)\n  subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)\n  return tree_unflatten(inner_treedef, subtrees)\n\ndef _num_leaves(treedef):\n  return 1 if treedef is leaf else sum(map(_num_leaves, treedef.children))\n\ndef _nested_treedef(inner, outer):\n  # just used in tree_transpose error checking\n  if outer is leaf:\n    return inner\n  else:\n    children = map(partial(_nested_treedef, inner), outer.children)\n    return PyTreeDef(outer.node_type, outer.node_data, tuple(children))\n\n\ndef tree_structure(tree):\n  _, spec = process_pytree(lambda _: None, tree)\n  return spec\n\n\nclass PyTreeDef(object):\n  def __init__(self, node_type, node_data, children):\n    self.node_type = node_type\n    self.node_data = node_data\n    self.children = children\n\n  def __repr__(self):\n    if self.node_data is None:\n      data_repr = \"\"\n    else:\n      data_repr = \"[{}]\".format(self.node_data)\n\n    return \"PyTree({}{}, [{}])\".format(self.node_type.name, data_repr,\n                                     ','.join(map(repr, self.children)))\n\n  def __hash__(self):\n    return hash((self.node_type, self.node_data, tuple(self.children)))\n\n  def __eq__(self, other):\n    if other is leaf:\n      return False\n    else:\n      return (self.node_type == other.node_type and\n              self.node_data == other.node_data and\n              self.children == other.children)\n\n  def __ne__(self, other):\n    return not self == other\n\n\nclass PyLeaf(object):\n  def __repr__(self):\n    return '*'\n\nleaf = PyLeaf()\n\ndef dict_to_iterable(xs):\n  keys = tuple(sorted(xs.keys()))\n  return tuple(map(xs.get, keys)), keys\n\nclass NodeType(object):\n  def __init__(self, name, to_iterable, from_iterable):\n    self.name = name\n    self.to_iterable = to_iterable\n    self.from_iterable = from_iterable\n\n  def __repr__(self):\n    return self.name\n\nnode_types = {}\n\ndef register_pytree_node(py_type, to_iterable, from_iterable):\n  assert py_type not in node_types\n  node_types[py_type] = NodeType(str(py_type), to_iterable, from_iterable)\n\nregister_pytree_node(tuple, lambda xs: (xs, None), lambda _, xs: tuple(xs))\nregister_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))\nregister_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))\nregister_pytree_node(type(None), lambda z: ((), None), lambda _, xs: None)\n\n\n# To handle namedtuples, we can't just use the standard table of node_types\n# because every namedtuple creates its own type and thus would require its own\n# entry in the table. Instead we use a heuristic check on the type itself to\n# decide whether it's a namedtuple type, and if so treat it as a pytree node.\ndef _get_node_type(maybe_tree):\n  t = type(maybe_tree)\n  return node_types.get(t) or _namedtuple_node(t)\n\ndef _namedtuple_node(t):\n  if t.__bases__ == (tuple,) and hasattr(t, '_fields'):\n    return NamedtupleNode\n\nNamedtupleNode = NodeType('namedtuple',\n                          lambda xs: (tuple(xs), type(xs)),\n                          lambda t, xs: t(*xs))\n", "path": "jax/tree_util.py"}]}