code: def tree_flatten(example, name="restructure"):
      """Returns a Bijector variant of tf.nest.flatten.
  
      To make it a Bijector, it has to know how to "unflatten" as
      well---unlike the real `tf.nest.flatten`, this can only flatten or
      unflatten a specific structure.  The `example` argument defines the
      structure.
  
      See also the `Restructure` bijector for general rearrangements.
  
      Args:
        example: A Tensor or (potentially nested) collection of Tensors.
        name: An optional Python string, inserted into names of TF ops
          created by this bijector.
  
      Returns:
        flatten: A Bijector whose `forward` method flattens structures
          parallel to `example` into a list of Tensors, and whose
          `inverse` method packs a list of Tensors of the right length
          into a structure parallel to `example`.
  
      #### Example
  
      ```python
      x = tf.constant(1)
      example = collections.OrderedDict([
          ('a', [x, x, x]),
          ('b', x)])
      bij = tfb.tree_flatten(example)
      ys = collections.OrderedDict([
          ('a', [1, 2, 3]),
          ('b', 4.)])
      bij.forward(ys)
      # Returns [1, 2, 3, 4.]
      ```
  
      """
      return invert.Invert(pack_sequence_as(example, name))
  
imports: import six
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: tree_flatten(example, name) and returns: invert