Make Layer track checkpoint dependencies
e.g. self.v = tf.Variable(1.) will now make the variable visible to tf.train.Checkpoint if "self" is a Layer subclass but not a Network. It should inherit from Module eventually, which will make it AutoCheckpointable anyway. But in the meantime this is causing confusion when subclassing Layer. Moves the tracking logic from Network to Layer. Fixes a few issues which came up, one bug where __radd__ was in the wrong order and one where collections.Counter was giving different results after wrapping, but otherwise just tangential issues with utilities inspecting _layer or special-casing dict instead of using collections.Mapping. PiperOrigin-RevId: 234825412
Loading
Please sign in to comment