Explicitly pass training argument to the layer in QAT module.

PiperOrigin-RevId: 370377067

Author: rino20

tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate.py

@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
+
 import tensorflow as tf
 
 deserialize_keras_object = tf.keras.utils.deserialize_keras_object
@@ -88,9 +90,11 @@ def __init__(self, layer, quantize_config=None, **kwargs):
       self._batch_input_shape = self.layer._batch_input_shape  # pylint: disable=protected-access
 
   def call(self, *args, **kwargs):
-    # TODO(b/185306646): Explicitly wants to pass training argument for the
-    # layer. Currently, we remove training argument.
-    if 'training' in kwargs:
+    arg = inspect.getfullargspec(self.layer.call).args
+
+    # Do not propagate the training bool to the underlying layer if it doesn't
+    # accepts the training bool.
+    if 'training' not in arg and 'training' in kwargs:
       del kwargs['training']
     return self.layer.call(*args, **kwargs)
 

tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate_test.py

@@ -19,9 +19,9 @@
 from __future__ import print_function
 
 import numpy as np
-
 import tensorflow as tf
 
+from tensorflow.python.keras.engine.base_layer import Layer
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_annotate
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_config as quantize_config_mod
 
@@ -52,6 +52,21 @@ def get_output_quantizers(self, layer):
     def get_config(self):
       return {}
 
+  def testAnnotateLayerCallPassesTraningBoolean(self):
+
+    class MockLayer(Layer):
+      self.training = None
+
+      def call(self, training=None):
+        self.training = training
+
+    layer = MockLayer()
+    wrapper = quantize_annotate.QuantizeAnnotate(layer=layer)
+    wrapper.call(training=True)
+    self.assertTrue(layer.training)
+    wrapper.call(training=False)
+    self.assertFalse(layer.training)
+
   def testAnnotatesKerasLayer(self):
     layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
     model = keras.Sequential([layer])