From 6aece71ebf756d32ea730576a7ff12d2cfc7b242 Mon Sep 17 00:00:00 2001 From: Thomas O'Malley Date: Tue, 26 May 2020 09:55:55 -0700 Subject: [PATCH] Reduce nest.map_structure and nest.pack_sequence_as time by ~20% for common use cases (nested structures of list, tuple, dict). Places relatively cheap type checks for list, tuple, and dict before other more expensive checks. Specifically, this avoids calling expensive checks like isinstance(structure, collections.abc.Mapping) and nest._is_named_tuple in the most common cases (since these abc isinstance checks take ~10x as long as normal isinstance checks). This reduces the Python overhead of a sample 10-layer Keras Functional Model __call__ by ~5%. PiperOrigin-RevId: 313211095 Change-Id: I227a3dc379eefef31060698d8c5be5f4bf2c1f50 --- tensorflow/python/util/nest.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/util/nest.py b/tensorflow/python/util/nest.py index 695cc4cc909..b4736bee142 100644 --- a/tensorflow/python/util/nest.py +++ b/tensorflow/python/util/nest.py @@ -215,7 +215,15 @@ def _yield_sorted_items(iterable): Yields: The iterable's (key, value) pairs, in order of sorted keys. """ - if isinstance(iterable, _collections_abc.Mapping): + # Ordered to check common structure types (list, tuple, dict) first. + if isinstance(iterable, list): + for item in enumerate(iterable): + yield item + # namedtuples handled separately to avoid expensive namedtuple check. + elif type(iterable) == tuple: # pylint: disable=unidiomatic-typecheck + for item in enumerate(iterable): + yield item + elif isinstance(iterable, (dict, _collections_abc.Mapping)): # Iterate through dictionaries in a deterministic order by sorting the # keys. Notice this means that we ignore the original order of `OrderedDict` # instances. This is intentional, to avoid potential bugs caused by mixing