Align default containers with jax (#14)

Author: janosg

.pre-commit-config.yaml

@@ -96,7 +96,7 @@ repos:
     hooks:
     -   id: interrogate
         args: [-v, --fail-under=20]
-        exclude: ^(docs|setup\.py)
+        exclude: ^(tests|docs|setup\.py)
 - repo: https://github.com/codespell-project/codespell
   rev: v2.1.0
   hooks:

src/pybaum/registry.py

@@ -1,7 +1,7 @@
 from pybaum.registry_entries import FUNC_DICT
 
 
-def get_registry(types=None, options=None, include_defaults=True):
+def get_registry(types=None, include_defaults=True):
     """Create a pytree registry.
 
     Args:
@@ -11,14 +11,15 @@ def get_registry(types=None, options=None, include_defaults=True):
             - "tuple"
             - "dict"
             - "list"
+            - :class:`collections.namedtuple` or :class:`typing.NamedTuple`
+            - :obj:`None`
+            - :class:`collections.OrderedDict`
             - "numpy.ndarray"
             - "pandas.Series"
             - "pandas.DataFrame"
-        options (dict): Option dictionary where the keys are names of types and the
-            values are keyword arguments that influence how containers are flattened
-            and unflattened.
         include_defaults (bool): Whether the default pytree containers "tuple", "dict"
-            and "list" should be included even if not specified in `types`.
+            "list", "None", "namedtuple" and "OrderedDict" should be included even if
+            not specified in `types`.
 
     Returns:
         dict: A pytree registry.
@@ -27,13 +28,12 @@ def get_registry(types=None, options=None, include_defaults=True):
     types = [] if types is None else types
 
     if include_defaults:
-        types = list(set(types) | {"list", "tuple", "dict"})
-
-    options = {} if options is None else options
+        default_types = {"list", "tuple", "dict", "None", "namedtuple", "OrderedDict"}
+        types = list(set(types) | default_types)
 
     registry = {}
     for typ in types:
-        new_entry = FUNC_DICT[typ](**options.get(typ, {}))
+        new_entry = FUNC_DICT[typ]()
         registry = {**registry, **new_entry}
 
     return registry

src/pybaum/registry_entries.py

@@ -1,5 +1,7 @@
 import itertools
-from functools import partial
+from collections import namedtuple
+from collections import OrderedDict
+from itertools import product
 
 from pybaum.config import IS_NUMPY_INSTALLED
 from pybaum.config import IS_PANDAS_INSTALLED
@@ -11,7 +13,20 @@
     import pandas as pd
 
 
+def _none():
+    """Create registry entry for NoneType."""
+    entry = {
+        type(None): {
+            "flatten": lambda tree: ([], None),  # noqa: U100
+            "unflatten": lambda aux_data, children: None,  # noqa: U100
+            "names": lambda tree: [],  # noqa: U100
+        }
+    }
+    return entry
+
+
 def _list():
+    """Create registry entry for list."""
     entry = {
         list: {
             "flatten": lambda tree: (tree, None),
@@ -23,6 +38,7 @@ def _list():
 
 
 def _dict():
+    """Create registry entry for dict."""
     entry = {
         dict: {
             "flatten": lambda tree: (list(tree.values()), list(tree)),
@@ -34,6 +50,7 @@ def _dict():
 
 
 def _tuple():
+    """Create registry entry for tuple."""
     entry = {
         tuple: {
             "flatten": lambda tree: (list(tree), None),
@@ -44,12 +61,41 @@ def _tuple():
     return entry
 
 
-def _numpy_array():
-    """Create a pytree declaration for numpy arrays.
+def _namedtuple():
+    """Create registry entry for namedtuple and NamedTuple."""
+    entry = {
+        namedtuple: {
+            "flatten": lambda tree: (list(tree), tree),
+            "unflatten": _unflatten_namedtuple,
+            "names": lambda tree: list(tree._fields),
+        },
+    }
+    return entry
+
+
+def _unflatten_namedtuple(aux_data, leaves):
+    replacements = dict(zip(aux_data._fields, leaves))
+    out = aux_data._replace(**replacements)
+    return out
 
-    To-Do: Add optional axis argument.
 
-    """
+def _ordereddict():
+    """Create registry entry for OrderedDict."""
+    entry = {
+        OrderedDict: {
+            "flatten": lambda tree: (list(tree.values()), list(tree)),
+            "unflatten": lambda aux_data, children: OrderedDict(
+                zip(aux_data, children)
+            ),
+            "names": lambda tree: list(map(str, list(tree))),
+        },
+    }
+    return entry
+
+
+def _numpy_array():
+    """Create registry entry for numpy.ndarray."""
+
     if IS_NUMPY_INSTALLED:
         entry = {
             np.ndarray: {
@@ -72,6 +118,7 @@ def _array_element_names(arr):
 
 
 def _pandas_series():
+    """Create registry entry for pandas.Series."""
     if IS_PANDAS_INSTALLED:
         entry = {
             pd.Series: {
@@ -88,69 +135,49 @@ def _pandas_series():
     return entry
 
 
-def _pandas_dataframe(columns=None):
+def _pandas_dataframe():
+    """Create registry entry for pandas.DataFrame."""
     if IS_PANDAS_INSTALLED:
         entry = {
             pd.DataFrame: {
-                "flatten": partial(_flatten_pandas_dataframe, columns=columns),
-                "unflatten": partial(_unflatten_pandas_dataframe),
-                "names": partial(_get_names_pandas_dataframe, columns=columns),
+                "flatten": _flatten_pandas_dataframe,
+                "unflatten": _unflatten_pandas_dataframe,
+                "names": _get_names_pandas_dataframe,
             }
         }
     else:
         entry = {}
     return entry
 
 
-def _flatten_pandas_dataframe(df, columns):
-    columns = _process_columns(df, columns)
-    flat = []
-    for col in columns:
-        flat += df[col].tolist()
-
-    aux_data = (columns, df.drop(columns=columns))
+def _flatten_pandas_dataframe(df):
+    flat = df.to_numpy().flatten().tolist()
+    aux_data = {"columns": df.columns, "index": df.index, "shape": df.shape}
     return flat, aux_data
 
 
 def _unflatten_pandas_dataframe(aux_data, leaves):
-    columns, empty_df = aux_data
-    out = empty_df.copy()
-    remaining_leaves = leaves
-    for col in columns:
-        out[col] = leaves[: len(empty_df)]
-        remaining_leaves = remaining_leaves[len(empty_df) :]
+    out = pd.DataFrame(
+        data=np.array(leaves).reshape(aux_data["shape"]),
+        columns=aux_data["columns"],
+        index=aux_data["index"],
+    )
     return out
 
 
-def _get_names_pandas_dataframe(df, columns):
-    columns = _process_columns(df, columns)
-    if len(columns) == 1:
-        out = list(df.index.map(_index_element_to_string))
-    else:
-        out = []
-        for col in df.columns:
-            out += list(df.index.map(partial(_index_element_to_string, prefix=col)))
+def _get_names_pandas_dataframe(df):
+    index_strings = list(df.index.map(_index_element_to_string))
+    out = ["_".join([loc, col]) for loc, col in product(index_strings, df.columns)]
     return out
 
 
-def _process_columns(df, columns):
-    if columns is None:
-        columns = df.columns
-    elif not isinstance(columns, list):
-        columns = [columns]
-    return columns
-
-
-def _index_element_to_string(element, prefix=None):
-    separator = "_"
+def _index_element_to_string(element):
     if isinstance(element, (tuple, list)):
         as_strings = [str(entry) for entry in element]
-        res_string = separator.join(as_strings)
+        res_string = "_".join(as_strings)
     else:
         res_string = str(element)
 
-    if prefix is not None:
-        res_string = separator.join([prefix, res_string])
     return res_string
 
 
@@ -161,4 +188,7 @@ def _index_element_to_string(element, prefix=None):
     "numpy.ndarray": _numpy_array,
     "pandas.Series": _pandas_series,
     "pandas.DataFrame": _pandas_dataframe,
+    "None": _none,
+    "namedtuple": _namedtuple,
+    "OrderedDict": _ordereddict,
 }

src/pybaum/tree_util.py

@@ -10,6 +10,7 @@
 
 from pybaum.equality import EQUALITY_CHECKERS
 from pybaum.registry import get_registry
+from pybaum.typecheck import get_type
 
 
 def tree_flatten(tree, is_leaf=None, registry=None):
@@ -80,14 +81,14 @@ def tree_just_flatten(tree, is_leaf=None, registry=None):
 
 def _tree_flatten(tree, is_leaf, registry):
     out = []
-    tree_type = type(tree)
+    tree_type = get_type(tree)
 
     if tree_type not in registry or is_leaf(tree):
         out.append(tree)
     else:
         subtrees, _ = registry[tree_type]["flatten"](tree)
         for subtree in subtrees:
-            if type(subtree) in registry:
+            if get_type(subtree) in registry:
                 out += _tree_flatten(subtree, is_leaf, registry)
             else:
                 out.append(subtree)
@@ -161,14 +162,14 @@ def tree_just_yield(tree, is_leaf=None, registry=None):
 
 def _tree_yield(tree, is_leaf, registry):
     out = []
-    tree_type = type(tree)
+    tree_type = get_type(tree)
 
     if tree_type not in registry or is_leaf(tree):
         yield tree
     else:
         subtrees, _ = registry[tree_type]["flatten"](tree)
         for subtree in subtrees:
-            if type(subtree) in registry:
+            if get_type(subtree) in registry:
                 yield from _tree_yield(subtree, is_leaf, registry)
             else:
                 yield subtree
@@ -211,15 +212,15 @@ def tree_unflatten(treedef, leaves, is_leaf=None, registry=None):
 
 def _tree_unflatten(treedef, leaves, is_leaf, registry):
     leaves = iter(leaves)
-    tree_type = type(treedef)
+    tree_type = get_type(treedef)
 
     if tree_type not in registry or is_leaf(treedef):
         return next(leaves)
     else:
         items, info = registry[tree_type]["flatten"](treedef)
         unflattened_items = []
         for item in items:
-            if type(item) in registry:
+            if get_type(item) in registry:
                 unflattened_items.append(
                     _tree_unflatten(item, leaves, is_leaf=is_leaf, registry=registry)
                 )
@@ -336,15 +337,15 @@ def leaf_names(tree, is_leaf=None, registry=None, separator="_"):
 
 def _leaf_names(tree, is_leaf, registry, separator, prefix=None):
     out = []
-    tree_type = type(tree)
+    tree_type = get_type(tree)
 
     if tree_type not in registry or is_leaf(tree):
         out.append(prefix)
     else:
         subtrees, info = registry[tree_type]["flatten"](tree)
         names = registry[tree_type]["names"](tree)
         for name, subtree in zip(names, subtrees):
-            if type(subtree) in registry:
+            if get_type(subtree) in registry:
                 out += _leaf_names(
                     subtree,
                     is_leaf=is_leaf,
@@ -424,7 +425,7 @@ def tree_equal(tree, other, is_leaf=None, registry=None, equality_checkers=None)
 
     if equal:
         for first, second in zip(first_flat, second_flat):
-            check_func = equality_checkers.get(type(first), lambda a, b: a == b)
+            check_func = equality_checkers.get(get_type(first), lambda a, b: a == b)
             equal = equal and check_func(first, second)
             if not equal:
                 break

src/pybaum/typecheck.py

@@ -0,0 +1,31 @@
+from collections import namedtuple
+
+
+def get_type(obj):
+    """namdetuple aware type check.
+
+    As in JAX we treat collections.namedtuple and typing.NamedTuple both as
+    namedtuple but the exact type is preserved in the unflatten function.
+
+    namedtuples are discovered by being instances of tuple and having a
+    ``_fields`` attribute as suggested by Raymond Hettinger
+    `here <https://bugs.python.org/issue7796>`_.
+
+    Moreover we check for the presence of a ``_replace`` method because we need when
+    unflattening pytrees.
+
+    This can produce false positives but in most cases would still result in desired
+    behavior.
+
+    Args:
+        obj: The object to be checked
+
+    Returns:
+        bool
+
+    """
+    if isinstance(obj, tuple) and hasattr(obj, "_fields") and hasattr(obj, "_replace"):
+        out = namedtuple
+    else:
+        out = type(obj)
+    return out