Replace already replaced layer in model_transformer

ModelTransformer had a subtle bug in which recursive
transforms weren't being applied in cases where candidate
layers existed. Else it worked fine.

Now replaced layers are added as candidate layers so
they can be replaced.

PiperOrigin-RevId: 321030782

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_layout_transform.py

@@ -65,4 +65,4 @@ def apply(self, model, layer_quantize_map):
     ]
     return model_transformer.ModelTransformer(
         model, transforms,
-        layer_quantize_map.keys(), layer_quantize_map).transform()
+        set(layer_quantize_map.keys()), layer_quantize_map).transform()

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer.py

@@ -212,9 +212,9 @@ def _match_layer_with_inputs(self, layer, pattern, is_head_node):
     # Inbound layers can have different order from the list of input patterns.
     # TODO(pulkitb): Fix by checking all permutations.
     input_match_layer_nodes = []
-    for input_layer, pattern in zip(input_layers, pattern.inputs):
+    for input_layer, pattern_ in zip(input_layers, pattern.inputs):
       match_layer_node = self._match_layer_with_inputs(
-          input_layer, pattern, is_head_node=False)
+          input_layer, pattern_, is_head_node=False)
       if not match_layer_node:
         return None
       input_match_layer_nodes.append(match_layer_node)
@@ -354,12 +354,13 @@ def _assign_inbounds_for_replacement(layer_node):
     def _add_replacement_layer(layer_node):
       """Recursively add new layers."""
       self._config['layers'].append(layer_node.layer)
+      layer_name = layer_node.layer['config']['name']
       if layer_node.weights:
-        self._layer_weights_map[layer_node.layer['config']
-                                ['name']] = layer_node.weights
+        self._layer_weights_map[layer_name] = layer_node.weights
       if layer_node.metadata:
-        self._layer_metadata_map[layer_node.layer['config']
-                                 ['name']] = layer_node.metadata
+        self._layer_metadata_map[layer_name] = layer_node.metadata
+      if self.candidate_layers:
+        self.candidate_layers.add(layer_name)
 
       for input_layer in layer_node.input_layers:
         _add_replacement_layer(input_layer)
@@ -399,12 +400,13 @@ def _add_replacement_nodes(first_layer_removed_index, replacement_nodes):
       i = first_layer_removed_index
       for replacement_node in replacement_nodes:
         self._config['layers'].insert(i, replacement_node.layer)
+        layer_name = replacement_node.layer['config']['name']
         if replacement_node.weights:
-          self._layer_weights_map[replacement_node.layer['config']
-                                  ['name']] = replacement_node.weights
+          self._layer_weights_map[layer_name] = replacement_node.weights
         if replacement_node.metadata:
-          self._layer_metadata_map[replacement_node.layer['config']
-                                   ['name']] = replacement_node.metadata
+          self._layer_metadata_map[layer_name] = replacement_node.metadata
+        if self.candidate_layers:
+          self.candidate_layers.add(layer_name)
         i += 1
 
     replacement_nodes = _get_replacement_nodes(replacement_layer_node)

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer_test.py

@@ -450,6 +450,39 @@ def testReplaceTreeOfLayers_WithTreeOfLayers(self):
     # TODO(pulkitb): Implement
     pass
 
+  @parameterized.parameters(['sequential', 'functional'])
+  def testReplace_AlreadyReplacedLayer_WithAnotherMatch(self, model_type):
+    """Verifies a layer replaced using a transform can be replaced again."""
+
+    class ReplaceReLUWithSoftmax(Transform):
+
+      def pattern(self):
+        return LayerPattern('ReLU')
+
+      def replacement(self, match_layer):
+        replace_layer = keras.layers.serialize(keras.layers.Softmax())
+        replace_layer['name'] = replace_layer['config']['name']
+        return LayerNode(replace_layer)
+
+    class ReplaceSoftmaxWithELU(Transform):
+
+      def pattern(self):
+        return LayerPattern('Softmax')
+
+      def replacement(self, match_layer):
+        replace_layer = keras.layers.serialize(keras.layers.ELU())
+        replace_layer['name'] = replace_layer['config']['name']
+        return LayerNode(replace_layer)
+
+    model = self._simple_dense_model(model_type)
+    transformed_model, _ = ModelTransformer(
+        model,
+        [ReplaceReLUWithSoftmax(), ReplaceSoftmaxWithELU()],
+        candidate_layers=set([layer.name for layer in model.layers])
+    ).transform()
+
+    self.assertEqual(transformed_model.layers[-1].__class__.__name__, 'ELU')
+
   @parameterized.parameters(['sequential', 'functional'])
   def testDoesNotMatchForever_IfReplacementEqualsMatch(self, model_type):
 

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/transforms.py

@@ -62,6 +62,12 @@ def __init__(self, class_name, config=None, inputs=None):
     self.config = config
     self.inputs = inputs
 
+  def __str__(self):
+    return '{} : {} <- [{}]'.format(
+        self.class_name,
+        self.config,
+        ', '.join([str(inp) for inp in self.inputs]))
+
 
 class LayerNode(object):
   """Represents a Node in a tree containing a layer.