Give a formal definition of nested structure.

PiperOrigin-RevId: 402654941
Change-Id: I287008c0e879d162354ed30cc61ad7b4453f273e
diff --git a/tensorflow/python/tools/api/generator/doc_srcs.py b/tensorflow/python/tools/api/generator/doc_srcs.py
index 3af5d0a..f2eef30 100644
--- a/tensorflow/python/tools/api/generator/doc_srcs.py
+++ b/tensorflow/python/tools/api/generator/doc_srcs.py
@@ -36,37 +36,66 @@
 
 
 _TENSORFLOW_DOC_SOURCES = {
-    'app': DocSource(docstring_module_name='platform.app'),
-    'bitwise': DocSource(docstring_module_name='ops.bitwise_ops'),
-    'compat': DocSource(docstring_module_name='util.compat'),
-    'distribute': DocSource(docstring_module_name='distribute.distribute_lib'),
-    'distributions': DocSource(
-        docstring_module_name='ops.distributions.distributions'),
-    'errors': DocSource(docstring_module_name='framework.errors'),
-    'experimental.numpy': DocSource(docstring_module_name='ops.numpy_ops'),
-    'gfile': DocSource(docstring_module_name='platform.gfile'),
-    'graph_util': DocSource(docstring_module_name='framework.graph_util'),
-    'image': DocSource(docstring_module_name='ops.image_ops'),
-    'linalg': DocSource(docstring_module_name='ops.linalg_ops'),
-    'logging': DocSource(docstring_module_name='ops.logging_ops'),
-    'losses': DocSource(docstring_module_name='ops.losses.losses'),
-    'manip': DocSource(docstring_module_name='ops.manip_ops'),
-    'math': DocSource(docstring_module_name='ops.math_ops'),
-    'metrics': DocSource(docstring_module_name='ops.metrics'),
-    'nn': DocSource(docstring_module_name='ops.nn_ops'),
-    'nn.rnn_cell': DocSource(docstring_module_name='ops.rnn_cell'),
-    'python_io': DocSource(docstring_module_name='lib.io.python_io'),
-    'ragged': DocSource(docstring_module_name='ops.ragged'),
-    'resource_loader': DocSource(
-        docstring_module_name='platform.resource_loader'),
-    'sets': DocSource(docstring_module_name='ops.sets'),
-    'signal': DocSource(docstring_module_name='ops.signal.signal'),
-    'sparse': DocSource(docstring_module_name='ops.sparse_ops'),
-    'strings': DocSource(docstring_module_name='ops.string_ops'),
-    'summary': DocSource(docstring_module_name='summary.summary'),
-    'sysconfig': DocSource(docstring_module_name='platform.sysconfig'),
-    'test': DocSource(docstring_module_name='platform.test'),
-    'train': DocSource(docstring_module_name='training.training'),
+    'app':
+        DocSource(docstring_module_name='platform.app'),
+    'bitwise':
+        DocSource(docstring_module_name='ops.bitwise_ops'),
+    'compat':
+        DocSource(docstring_module_name='util.compat'),
+    'distribute':
+        DocSource(docstring_module_name='distribute.distribute_lib'),
+    'distributions':
+        DocSource(docstring_module_name='ops.distributions.distributions'),
+    'errors':
+        DocSource(docstring_module_name='framework.errors'),
+    'experimental.numpy':
+        DocSource(docstring_module_name='ops.numpy_ops'),
+    'gfile':
+        DocSource(docstring_module_name='platform.gfile'),
+    'graph_util':
+        DocSource(docstring_module_name='framework.graph_util'),
+    'image':
+        DocSource(docstring_module_name='ops.image_ops'),
+    'linalg':
+        DocSource(docstring_module_name='ops.linalg_ops'),
+    'logging':
+        DocSource(docstring_module_name='ops.logging_ops'),
+    'losses':
+        DocSource(docstring_module_name='ops.losses.losses'),
+    'manip':
+        DocSource(docstring_module_name='ops.manip_ops'),
+    'math':
+        DocSource(docstring_module_name='ops.math_ops'),
+    'metrics':
+        DocSource(docstring_module_name='ops.metrics'),
+    'nest':
+        DocSource(docstring_module_name='util.nest'),
+    'nn':
+        DocSource(docstring_module_name='ops.nn_ops'),
+    'nn.rnn_cell':
+        DocSource(docstring_module_name='ops.rnn_cell'),
+    'python_io':
+        DocSource(docstring_module_name='lib.io.python_io'),
+    'ragged':
+        DocSource(docstring_module_name='ops.ragged'),
+    'resource_loader':
+        DocSource(docstring_module_name='platform.resource_loader'),
+    'sets':
+        DocSource(docstring_module_name='ops.sets'),
+    'signal':
+        DocSource(docstring_module_name='ops.signal.signal'),
+    'sparse':
+        DocSource(docstring_module_name='ops.sparse_ops'),
+    'strings':
+        DocSource(docstring_module_name='ops.string_ops'),
+    'summary':
+        DocSource(docstring_module_name='summary.summary'),
+    'sysconfig':
+        DocSource(docstring_module_name='platform.sysconfig'),
+    'test':
+        DocSource(docstring_module_name='platform.test'),
+    'train':
+        DocSource(docstring_module_name='training.training'),
 }
 
 _ESTIMATOR_DOC_SOURCES = {
diff --git a/tensorflow/python/util/nest.py b/tensorflow/python/util/nest.py
index d39c1f6..683dbc6 100644
--- a/tensorflow/python/util/nest.py
+++ b/tensorflow/python/util/nest.py
@@ -13,29 +13,74 @@
 # limitations under the License.
 # ==============================================================================
 
-"""## Functions for working with arbitrarily nested sequences of elements.
+"""Functions that work with structures.
 
-This module can perform operations on nested structures. A nested structure is a
-Python collection that can contain further collections as well as other objects
-called atoms. Note that numpy arrays are considered atoms.
+A structure is either:
 
-nest recognizes the following types of collections:
-  1.tuple
-  2.namedtuple
-  3.dict
-  4.orderedDict
-  5.MutableMapping
-  6.attr.s
+* one of the recognized Python collections, holding _nested structures_;
+* a value of any other type, typically a TensorFlow data type like Tensor,
+  Variable, or of compatible types such as int, float, ndarray, etc. these are
+  commonly referred to as _atoms_ of the structure.
 
-attr.s decorated classes (http://www.attrs.org) are also supported, in the
-same way as `namedtuple`.
+A structure of type `T` is a structure whose atomic items are of type `T`.
+For example, a structure of `tf.Tensor` only contains `tf.Tensor` as its atoms.
 
-The utilities here assume (and do not check) that the nested structures form a
-'tree', i.e., no references in the structure of the input of these functions
-should be recursive.
+Historically a _nested structure_ was called a _nested sequence_ in TensorFlow.
+A nested structure is sometimes called a _nest_ or a _tree_, but the formal
+name _nested structure_ is preferred.
 
-Example structures: `((3, 4), 5, (6, 7, (9, 10), 8))`, `(np.array(0),
-  (np.array([3, 4]), tf.constant([3, 4])))`
+Refer to [Nesting Data Structures]
+(https://en.wikipedia.org/wiki/Nesting_(computing)#Data_structures).
+
+The following collection types are recognized by `tf.nest` as nested
+structures:
+
+* `collections.abc.Sequence` (except `string` and `bytes`).
+  This includes `list`, `tuple`, and `namedtuple`.
+* `collections.abc.Mapping` (with sortable keys).
+  This includes `dict` and `collections.OrderedDict`.
+* `collections.abc.MappingView` (with sortable keys).
+* [`attr.s` classes](https://www.attrs.org/).
+
+Any other values are considered **atoms**.  Not all collection types are
+considered nested structures.  For example, the following types are
+considered atoms:
+
+* `set`
+* [`dataclass` classes](https://docs.python.org/library/dataclasses.html)
+* `tf.Tensor`
+* `numpy.array`
+
+`tf.nest.is_nested` checks whether an object is a nested structure or an atom.
+For example:
+
+  >>> tf.nest.is_nested("1234")
+  False
+  >>> tf.nest.is_nested([1, 3, [4, 5]])
+  True
+  >>> tf.nest.is_nested(((7, 8), (5, 6)))
+  True
+  >>> tf.nest.is_nested([])
+  True
+  >>> tf.nest.is_nested({"a": 1, "b": 2})
+  True
+  >>> tf.nest.is_nested({"a": 1, "b": 2}.keys())
+  True
+  >>> tf.nest.is_nested({"a": 1, "b": 2}.values())
+  True
+  >>> tf.nest.is_nested({"a": 1, "b": 2}.items())
+  True
+  >>> tf.nest.is_nested(set([1, 2]))
+  False
+  >>> ones = tf.ones([2, 3])
+  >>> tf.nest.is_nested(ones)
+  False
+
+Note: A proper structure shall form a tree. The user shall ensure there is no
+cyclic references within the items in the structure,
+i.e., no references in the structure of the input of these functions
+should be recursive. The behavior is undefined if there is a cycle.
+
 """
 
 import collections as _collections
@@ -65,7 +110,7 @@
 )
 
 _INPUT_TREE_SMALLER_THAN_SHALLOW_TREE = (
-    "The input_tree has fewer elements than the shallow_tree. Input structure "
+    "The input_tree has fewer items than the shallow_tree. Input structure "
     "has length {input_size}, while shallow structure has length "
     "{shallow_size}.")
 
@@ -144,7 +189,7 @@
     instance: an instance of `tuple`, `list`, `namedtuple`, `dict`,
         `collections.OrderedDict`, or `composite_tensor.Composite_Tensor`
         or `type_spec.TypeSpec`.
-    args: elements to be converted to the `instance` type.
+    args: items to be converted to the `instance` type.
 
   Returns:
     `args` with the type of `instance`.
@@ -270,60 +315,10 @@
 
 @tf_export("nest.is_nested")
 def is_nested(seq):
-  """Returns true if its input is a nested structure (or just "nest").
+  """Returns true if its input is a nested structure.
 
-  In particular, the following collection types are recognized by `tf.nest`
-  as nested structures:
-
-  * `collections.abc.Sequence` (except `string` and `bytes`).
-    This includes `list`, `tuple`, and `namedtuple`.
-  * `collections.abc.Mapping` (with sortable keys).
-    This includes `dict` and `collections.OrderedDict`.
-  * `collections.abc.MappingView` (with sortable keys).
-  * [`attr.s` classes](https://www.attrs.org/).
-
-  Any other values are considered **atoms**.  Not all collection types are
-  considered nested structures.  For example, the following types are
-  considered atoms:
-
-    * `set`
-    * [`dataclass` classes](https://docs.python.org/library/dataclasses.html)
-    * `tf.Tensor`
-    * `numpy.array`
-
-  >>> tf.nest.is_nested("1234")
-  False
-
-  >>> tf.nest.is_nested([1, 3, [4, 5]])
-  True
-
-  >>> tf.nest.is_nested(((7, 8), (5, 6)))
-  True
-
-  >>> tf.nest.is_nested([])
-  True
-
-  >>> tf.nest.is_nested({"a": 1, "b": 2})
-  True
-
-  >>> tf.nest.is_nested({"a": 1, "b": 2}.keys())
-  True
-
-  >>> tf.nest.is_nested({"a": 1, "b": 2}.values())
-  True
-
-  >>> tf.nest.is_nested({"a": 1, "b": 2}.items())
-  True
-
-  >>> tf.nest.is_nested(set([1, 2]))
-  False
-
-  >>> ones = tf.ones([2, 3])
-  >>> tf.nest.is_nested(ones)
-  False
-
-  Note: nested structures have sometimes been referred to as "nested sequences,"
-  but "nested structures" and "nests" are the preferred terms.
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a nested structure.
 
   Args:
     seq: the value to test.
@@ -336,11 +331,12 @@
 
 @tf_export("nest.flatten")
 def flatten(structure, expand_composites=False):
-  """Returns a flat list from a given nested structure.
+  """Returns a flat list from a given structure.
 
-  If nest is not a structure , tuple (or a namedtuple), dict, or an attrs class,
-  then returns a single-element list:
-    [nest].
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
+
+  If the structure is an atom, then returns a single-item list: [structure].
 
   This is the inverse of the `nest.pack_sequence_as` method that takes in a
   flattened list and re-packs it into the nested structure.
@@ -416,8 +412,8 @@
      <tf.Tensor: shape=(4,), dtype=int64, numpy=array([0, 4, 4, 7])>]
 
   Args:
-    structure: an arbitrarily nested structure. Note, numpy arrays are
-      considered atoms and are not flattened.
+    structure: an atom or a nested structure. Note, numpy arrays are considered
+      atoms and are not flattened.
     expand_composites: If true, then composite tensors such as
       `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
       component tensors.
@@ -458,16 +454,19 @@
                           expand_composites=False):
   """Asserts that two structures are nested in the same way.
 
-  Note the method does not check the types of data inside the structures.
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
+
+  Note the method does not check the types of atoms inside the structures.
 
   Examples:
 
-  * These scalar vs. scalar comparisons will pass:
+  * These atom vs. atom comparisons will pass:
 
     >>> tf.nest.assert_same_structure(1.5, tf.Variable(1, tf.uint32))
     >>> tf.nest.assert_same_structure("abc", np.array([1, 2]))
 
-  * These sequence vs. sequence comparisons will pass:
+  * These nested structure vs. nested structure comparisons will pass:
 
     >>> structure1 = (((1, 2), 3), 4, (5, 6))
     >>> structure2 = ((("foo1", "foo2"), "foo3"), "foo4", ("foo5", "foo6"))
@@ -516,22 +515,23 @@
     TypeError: The two structures don't have the same nested structure
 
   Args:
-    nest1: an arbitrarily nested structure.
-    nest2: an arbitrarily nested structure.
-    check_types: if `True` (default) types of sequences are checked as well,
-      including the keys of dictionaries. If set to `False`, for example a
-      list and a tuple of objects will look the same if they have the same
-      size. Note that namedtuples with identical name and fields are always
-      considered to have the same shallow structure. Two types will also be
-      considered the same if they are both list subtypes (which allows "list"
-      and "_ListWrapper" from trackable dependency tracking to compare
-      equal).
+    nest1: an atom or a nested structure.
+    nest2: an atom or a nested structure.
+    check_types: if `True` (default) types of structures are checked as well,
+      including the keys of dictionaries. If set to `False`, for example a list
+      and a tuple of objects will look the same if they have the same size. Note
+      that namedtuples with identical name and fields are always considered to
+      have the same shallow structure. Two types will also be considered the
+      same if they are both list subtypes (which allows "list" and
+      "_ListWrapper" from trackable dependency tracking to compare equal).
+      `check_types=True` only checks type of sub-structures. The types of atoms
+      are not checked.
     expand_composites: If true, then composite tensors such as
       `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
       component tensors.
 
   Raises:
-    ValueError: If the two structures do not have the same number of elements or
+    ValueError: If the two structures do not have the same number of atoms or
       if the two structures are not nested in the same way.
     TypeError: If the two structures differ in the type of sequence in any of
       their substructures. Only possible if `check_types` is `True`.
@@ -589,7 +589,7 @@
   """Helper function for pack_sequence_as.
 
   Args:
-    structure: Substructure (list / tuple / dict) to mimic.
+    structure: structure to mimic.
     flat: Flattened values to output substructure for.
     index: Index at which to start reading from flat.
     is_seq: Function used to test if a value should be treated as a sequence.
@@ -603,7 +603,7 @@
                  format.
 
   Raises:
-    ValueError: if `structure` contains more elements than `flat`
+    ValueError: if `structure` contains more atoms than `flat`
       (assuming indexing starts from `index`).
   """
   packed = []
@@ -631,15 +631,15 @@
 
   if not is_seq(flat_sequence):
     raise TypeError(
-        "Attempted to pack value:\n  {}\ninto a sequence, but found "
-        "incompatible type `{}` instead."
-        .format(truncate(flat_sequence, 100), type(flat_sequence)))
+        "Attempted to pack value:\n  {}\ninto a structure, but found "
+        "incompatible type `{}` instead.".format(
+            truncate(flat_sequence, 100), type(flat_sequence)))
 
   if not is_seq(structure):
     if len(flat_sequence) != 1:
       raise ValueError(
           "The target structure is of type `{}`\n  {}\nHowever the input "
-          "structure is a sequence ({}) of length {}.\n  {}\nnest cannot "
+          "is a sequence ({}) of length {}.\n  {}\nnest cannot "
           "guarantee that it is safe to map one to the other.".format(
               type(structure), truncate(structure, 100), type(flat_sequence),
               len(flat_sequence), truncate(flat_sequence, 100)))
@@ -654,8 +654,8 @@
     flat_structure = flatten(structure, expand_composites=expand_composites)
     if len(flat_structure) != len(flat_sequence):
       raise ValueError(
-          "Could not pack sequence. Structure had %d elements, but "
-          "flat_sequence had %d elements.  Structure: %s, flat_sequence: %s." %
+          "Could not pack sequence. Structure had %d atoms, but "
+          "flat_sequence had %d items.  Structure: %s, flat_sequence: %s." %
           (len(flat_structure), len(flat_sequence), structure, flat_sequence))
   return sequence_fn(structure, packed)
 
@@ -664,7 +664,10 @@
 def pack_sequence_as(structure, flat_sequence, expand_composites=False):
   """Returns a given flattened sequence packed into a given structure.
 
-  If `structure` is a scalar, `flat_sequence` must be a single-element list;
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
+
+  If `structure` is an atom, `flat_sequence` must be a single-item list;
   in this case the return value is `flat_sequence[0]`.
 
   If `structure` is or contains a dict instance, the keys will be sorted to
@@ -768,7 +771,7 @@
 
   Raises:
     ValueError: If `flat_sequence` and `structure` have different
-      element counts.
+      atom counts.
     TypeError: `structure` is or contains a dict with non-sortable keys.
   """
   return _pack_sequence_as(structure, flat_sequence, expand_composites)
@@ -776,10 +779,13 @@
 
 @tf_export("nest.map_structure")
 def map_structure(func, *structure, **kwargs):
-  """Applies `func` to each entry in `structure` and returns a new structure.
+  """Creates a new structure by applying `func` to each atom in `structure`.
 
-  Applies `func(x[0], x[1], ...)` where x[i] is an entry in
-  `structure[i]`.  All structures in `structure` must have the same arity,
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
+
+  Applies `func(x[0], x[1], ...)` where x[i] enumerates all atoms in
+  `structure[i]`.  All items in `structure` must have the same arity,
   and the return value will contain results with the same structure layout.
 
   Examples:
@@ -831,28 +837,24 @@
 
   Args:
     func: A callable that accepts as many arguments as there are structures.
-    *structure: scalar, or tuple or dict or list of constructed scalars and/or
-      other tuples/lists, or scalars.  Note: numpy arrays are considered as
-      scalars.
+    *structure: atom or nested structure.
     **kwargs: Valid keyword args are:
-
-      * `check_types`: If set to `True` (default) the types of
-        iterables within the structures have to be same (e.g.
-        `map_structure(func, [1], (1,))` raises a `TypeError`
-        exception). To allow this set this argument to `False`.
-        Note that namedtuples with identical name and fields are always
+      * `check_types`: If set to `True` (default) the types of iterables within
+        the structures have to be same (e.g. `map_structure(func, [1], (1,))`
+        raises a `TypeError` exception). To allow this set this argument to
+        `False`. Note that namedtuples with identical name and fields are always
         considered to have the same shallow structure.
-      * `expand_composites`: If set to `True`, then composite tensors such
-        as `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into
-        their component tensors.  If `False` (the default), then composite
-        tensors are not expanded.
+      * `expand_composites`: If set to `True`, then composite tensors such as
+        `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
+        component tensors.  If `False` (the default), then composite tensors are
+        not expanded.
 
   Returns:
-    A new structure with the same arity as `structure`, whose values correspond
-    to `func(x[0], x[1], ...)` where `x[i]` is a value in the corresponding
-    location in `structure[i]`. If there are different sequence types and
-    `check_types` is `False` the sequence types of the first structure will be
-    used.
+    A new structure with the same arity as `structure[0]`, whose atoms
+    correspond to `func(x[0], x[1], ...)` where `x[i]` is the atom in the
+    corresponding location in `structure[i]`. If there are different structure
+    types and `check_types` is `False` the structure types of the first
+    structure will be used.
 
   Raises:
     TypeError: If `func` is not callable or if the structures do not match
@@ -1131,15 +1133,18 @@
                   expand_composites=False):
   """Flattens `input_tree` up to `shallow_tree`.
 
-  Any further depth in structure in `input_tree` is retained as elements in the
-  partially flatten output.
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
 
-  If `shallow_tree` and `input_tree` are not sequences, this returns a
-  single-element list: `[input_tree]`.
+  Any further depth in structure in `input_tree` is retained as structures in
+  the partially flatten output.
+
+  If `shallow_tree` and `input_tree` are atoms, this returns a
+  single-item list: `[input_tree]`.
 
   Use Case:
 
-  Sometimes we may wish to partially flatten a nested sequence, retaining some
+  Sometimes we may wish to partially flatten a structure, retaining some
   of the nested structure. We achieve this by specifying a shallow structure,
   `shallow_tree`, we wish to flatten up to.
 
@@ -1172,7 +1177,7 @@
   # ['a', 1, 'b', 2, 'c', 3, 'd', 4]
   ```
 
-  Non-Sequence Edge Cases:
+  Edge Cases for atoms:
 
   ```python
   flatten_up_to(0, 0)  # Output: [0]
@@ -1183,8 +1188,8 @@
 
   Args:
     shallow_tree: a possibly pruned structure of input_tree.
-    input_tree: an arbitrarily nested structure or a scalar object.
-      Note, numpy arrays are considered scalars.
+    input_tree: an atom or a nested structure.
+      Note, numpy arrays are considered atoms.
     check_types: bool. If True, check that each node in shallow_tree has the
       same type as the corresponding node in input_tree.
     expand_composites: If true, then composite tensors such as
@@ -1196,10 +1201,10 @@
     the structure of `shallow_tree`.
 
   Raises:
-    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
-    TypeError: If the sequence types of `shallow_tree` are different from
+    TypeError: If `shallow_tree` is a nested structure but `input_tree` is not.
+    TypeError: If the structure types of `shallow_tree` are different from
       `input_tree`.
-    ValueError: If the sequence lengths of `shallow_tree` are different from
+    ValueError: If the structure lengths of `shallow_tree` are different from
       `input_tree`.
   """
   is_seq = is_sequence_or_composite if expand_composites else is_sequence
@@ -1217,14 +1222,14 @@
                                    expand_composites=False):
   """Flattens `input_tree` up to `shallow_tree`.
 
-  Any further depth in structure in `input_tree` is retained as elements in the
-  partially flattened output.
+  Any further depth in structure in `input_tree` is retained as structures in
+  the partially flattened output.
 
   Returns a list of (path, value) pairs, where value a leaf node in the
   flattened tree, and path is the tuple path of that leaf in input_tree.
 
   If `shallow_tree` and `input_tree` are not sequences, this returns a
-  single-element list: `[((), input_tree)]`.
+  single-item list: `[((), input_tree)]`.
 
   Use Case:
 
@@ -1288,8 +1293,8 @@
 
   Args:
     shallow_tree: a possibly pruned structure of input_tree.
-    input_tree: an arbitrarily nested structure or a scalar object.
-      Note, numpy arrays are considered scalars.
+    input_tree: an atom or a nested structure.
+      Note, numpy arrays are considered atoms.
     check_types: bool. If True, check that each node in shallow_tree has the
       same type as the corresponding node in input_tree.
     expand_composites: If true, then composite tensors such as
@@ -1301,10 +1306,10 @@
     the structure of `shallow_tree`.
 
   Raises:
-    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
-    TypeError: If the sequence types of `shallow_tree` are different from
+    TypeError: If `shallow_tree` is a nested structure but `input_tree` is not.
+    TypeError: If the structure types of `shallow_tree` are different from
       `input_tree`.
-    ValueError: If the sequence lengths of `shallow_tree` are different from
+    ValueError: If the structure lengths of `shallow_tree` are different from
       `input_tree`.
   """
   is_seq = is_sequence_or_composite if expand_composites else is_sequence
@@ -1324,7 +1329,7 @@
   Use Case:
 
   Sometimes we wish to apply a function to a partially flattened
-  sequence (for example when the function itself takes sequence inputs). We
+  structure (for example when the function itself takes structure inputs). We
   achieve this by specifying a shallow structure, `shallow_tree` we wish to
   flatten up to.
 
@@ -1367,12 +1372,11 @@
   ```
 
   Args:
-    shallow_tree: a shallow tree, common to all the inputs.
+    shallow_tree: a shallow structure, common to all the inputs.
     func: callable which will be applied to each input individually.
-    *inputs: arbitrarily nested combination of objects that are compatible with
-        shallow_tree. The function `func` is applied to corresponding
-        partially flattened elements of each input, so the function must support
-        arity of `len(inputs)`.
+    *inputs: structures that are compatible with shallow_tree. The function
+        `func` is applied to corresponding structures due to partial flattening
+        of each input, so the function must support arity of `len(inputs)`.
     **kwargs: kwargs to feed to func(). Special kwarg
       `check_types` is not passed to func, but instead determines whether the
       types of iterables within the structures have to be same (e.g.
@@ -1380,10 +1384,10 @@
       this set this argument to `False`.
 
   Raises:
-    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
-    TypeError: If the sequence types of `shallow_tree` are different from
+    TypeError: If `shallow_tree` is a nested structure but `input_tree` is not.
+    TypeError: If the structure types of `shallow_tree` are different from
       `input_tree`.
-    ValueError: If the sequence lengths of `shallow_tree` are different from
+    ValueError: If the structure lengths of `shallow_tree` are different from
       `input_tree`.
 
   Returns:
@@ -1441,11 +1445,11 @@
   ```
 
   Args:
-    shallow_tree: a shallow tree, common to all the inputs.
+    shallow_tree: a shallow structure, common to all the inputs.
     func: callable that takes args (path, inputs_0_value, ... , inputs_N_value),
-      where path is a tuple path to a leaf node in shallow_tree, and
+      where path is a tuple path to an atom in shallow_tree, and
       inputs_i_value is the corresponding value from inputs[i].
-    *inputs: nested structures that are all structurally compatible with
+    *inputs: structures that are all structurally compatible with
         shallow_tree.
     **kwargs: kwargs to feed to func(). Special kwarg
       `check_types` is not passed to func, but instead determines whether the
@@ -1454,10 +1458,11 @@
       this set this argument to `False`.
 
   Raises:
-    TypeError: If `shallow_tree` is a sequence but one of `*inputs` is not.
-    TypeError: If the sequence types of `shallow_tree` are different from
+    TypeError: If `shallow_tree` is a nested structure but one of `*inputs` is
+      not.
+    TypeError: If the structure types of `shallow_tree` are different from
       `input_tree`.
-    ValueError: If the sequence lengths of `shallow_tree` are different from
+    ValueError: If the structure lengths of `shallow_tree` are different from
       `input_tree`.
 
   Returns:
@@ -1478,7 +1483,7 @@
         check_types=check_types,
         expand_composites=expand_composites)
 
-  # Flatten each input separately, apply the function to corresponding elements,
+  # Flatten each input separately, apply the function to corresponding items,
   # then repack based on the structure of the first input.
   flat_value_gen = (
       flatten_up_to(  # pylint: disable=g-complex-comprehension
@@ -1522,8 +1527,8 @@
     `map_structure_up_to` and `flatten_up_to`.
 
   Raises:
-    TypeError: if `traverse_fn` returns a sequence for a non-sequence input,
-      or a structure with depth higher than 1 for a sequence input,
+    TypeError: if `traverse_fn` returns a nested structure for an atom input.
+      or a structure with depth higher than 1 for a nested structure input,
       or if any leaf values in the returned structure or scalar are not type
       `bool`.
   """
@@ -1570,6 +1575,9 @@
 def yield_flat_paths(nest, expand_composites=False):
   """Yields paths for some nested structure.
 
+  Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
+  for the definition of a structure.
+
   Paths are lists of objects which can be str-converted, which may include
   integers or other types which are used as indices in a dict.
 
@@ -1610,11 +1618,11 @@
 
 def flatten_with_joined_string_paths(structure, separator="/",
                                      expand_composites=False):
-  """Returns a list of (string path, data element) tuples.
+  """Returns a list of (string path, atom) tuples.
 
   The order of tuples produced matches that of `nest.flatten`. This allows you
   to flatten a nested structure while keeping information about where in the
-  structure each data element was located. See `nest.yield_flat_paths`
+  structure each atom was located. See `nest.yield_flat_paths`
   for more information.
 
   Args:
@@ -1626,7 +1634,7 @@
       component tensors.
 
   Returns:
-    A list of (string, data element) tuples.
+    A list of (string, atom) tuples.
   """
   flat_paths = yield_flat_paths(structure, expand_composites=expand_composites)
   def stringify_and_join(path_elements):
@@ -1638,11 +1646,11 @@
 
 
 def flatten_with_tuple_paths(structure, expand_composites=False):
-  """Returns a list of `(tuple_path, leaf_element)` tuples.
+  """Returns a list of `(tuple_path, atom)` tuples.
 
   The order of pairs produced matches that of `nest.flatten`. This allows you
   to flatten a nested structure while keeping information about where in the
-  structure each data element was located. See `nest.yield_flat_paths`
+  structure each atom was located. See `nest.yield_flat_paths`
   for more information about tuple paths.
 
   Args:
@@ -1652,9 +1660,9 @@
       component tensors.
 
   Returns:
-    A list of `(tuple_path, leaf_element)` tuples. Each `tuple_path` is a tuple
+    A list of `(tuple_path, atom)` tuples. Each `tuple_path` is a tuple
     of indices and/or dictionary keys that uniquely specify the path to
-    `leaf_element` within `structure`.
+    `atom` within `structure`.
   """
   return list(zip(yield_flat_paths(structure,
                                    expand_composites=expand_composites),
@@ -1665,7 +1673,7 @@
 def list_to_tuple(structure):
   """Replace all lists with tuples.
 
-  The fork of nest that tf.data uses treats lists as single elements, while
+  The fork of nest that tf.data uses treats lists as atoms, while
   tf.nest treats them as structures to recurse into. Keras has chosen to adopt
   the latter convention, and must therefore deeply replace all lists with tuples
   before passing structures to Dataset.from_generator.
diff --git a/tensorflow/python/util/nest_test.py b/tensorflow/python/util/nest_test.py
index 19e2304..234ef62 100644
--- a/tensorflow/python/util/nest_test.py
+++ b/tensorflow/python/util/nest_test.py
@@ -72,7 +72,7 @@
   PointXY = collections.namedtuple("Point", ["x", "y"])  # pylint: disable=invalid-name
   unsafe_map_pattern = ("nest cannot guarantee that it is safe to map one to "
                         "the other.")
-  bad_pack_pattern = ("Attempted to pack value:\n  .+\ninto a sequence, but "
+  bad_pack_pattern = ("Attempted to pack value:\n  .+\ninto a structure, but "
                       "found incompatible type `<(type|class) 'str'>` instead.")
 
   if attr:
@@ -279,8 +279,7 @@
 
   def testPackSequenceAs_wrongLengthsError(self):
     with self.assertRaisesRegex(
-        ValueError,
-        "Structure had 2 elements, but flat_sequence had 3 elements."):
+        ValueError, "Structure had 2 atoms, but flat_sequence had 3 items."):
       nest.pack_sequence_as(["hello", "world"],
                             ["and", "goodbye", "again"])
 
@@ -288,8 +287,7 @@
     val = ragged_tensor.RaggedTensor.from_row_splits(values=[1],
                                                      row_splits=[0, 1])
     with self.assertRaisesRegex(
-        ValueError,
-        "Structure had 2 elements, but flat_sequence had 1 elements."):
+        ValueError, "Structure had 2 atoms, but flat_sequence had 1 items."):
       nest.pack_sequence_as(val, [val], expand_composites=True)
 
   @test_util.assert_no_new_pyobjects_executing_eagerly