commit d2664949539e69aadd5050329104b4a81a217ee3
author Thomas O'Malley <omalleyt@google.com> Fri Oct 30 12:22:31 2020 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Oct 30 12:27:12 2020 -0700
tree db497ae733e217df3562ea3a12023fcde615fa66
parent 9c3c3855bd8404ce6c9191ee8043b080e7b3ac9b

Support TF Modules inside Keras Layers and Models.

With this change, it is now possible to mix-and-match tf.keras.Layers and
tf.Modules inside a tf.keras.Model and everything will be tracked properly.

- Variables in tf.Modules that are set as attributes of custom Layers and
  Models now show up properly in properties such as Layer.trainable_variables
  and Model.trainable_variables.
- tf.Modules do not show up in Model.layers. Instead, a new method
  Layer._flatten_modules is added that iterates over tf.Modules and Layers in
  the order that Keras expects. The existing method Layer.submodules (inherited
  from tf.Module) can still be used to iterate over tf.Modules and Layer with the
  tf.Module ordering. Layer._flatten_layers is built on top of
  Layer._flatten_modules
- Layer._layers is renamed to Layer._self_tracked_trackables to avoid naming
  conflicts with user-defined attributes (and to reflect that this attr
  contains Layers, Modules, and TrackableDataStructures)
- A new property is added to tf.Module to enable this, namely
  tf.Module.non_trainable_variables

PiperOrigin-RevId: 339917644
Change-Id: I96a7302745280a6261de8c4295c5cbf5f4d7dd5c