From 63ae3830c562108ec07f86b978993bcc040e5547 Mon Sep 17 00:00:00 2001 From: Ricardo Torres Date: Fri, 12 Jun 2026 10:50:00 -0700 Subject: [PATCH] Fix an issue where multiple instances of DefaultNBitQuantizeRegistry could interfere with each other's configuration. PiperOrigin-RevId: 931228650 --- .../default_n_bit/default_n_bit_quantize_registry.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py b/tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py index d33dc67be..b7bf8c569 100644 --- a/tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py +++ b/tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py @@ -185,9 +185,15 @@ def __init__(self, disable_per_axis=False, self._num_bits_activation = num_bits_activation self._layer_quantize_map = {} for quantize_info in self._LAYER_QUANTIZE_INFO: - quantize_info.num_bits_weight = num_bits_weight - quantize_info.num_bits_activation = num_bits_activation - self._layer_quantize_map[quantize_info.layer_type] = quantize_info + new_quantize_info = _QuantizeInfo( + layer_type=quantize_info.layer_type, + weight_attrs=quantize_info.weight_attrs, + activation_attrs=quantize_info.activation_attrs, + quantize_output=quantize_info.quantize_output, + num_bits_weight=num_bits_weight, + num_bits_activation=num_bits_activation, + ) + self._layer_quantize_map[new_quantize_info.layer_type] = new_quantize_info # Hack for `Activation` layer. That is the only layer with a separate # QuantizeConfig.