Supports custom quantizeWrapper instead of quantize annotation.

PiperOrigin-RevId: 377240567

Author: Xhark

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

@@ -399,8 +399,10 @@ def from_config(cls, config):
       return cls(**config)
 
   def _quantize(layer):  # pylint: disable=missing-docstring
-    if (layer.name not in layer_quantize_map and
-        layer.name not in requires_output_quantize):
+    if ((layer.name not in layer_quantize_map and
+         layer.name not in requires_output_quantize) or
+        (isinstance(layer, quantize_wrapper.QuantizeWrapper))):
+      # It supports for custom QuantizeWrapper.
       return layer
 
     if layer.name in requires_output_quantize: