Activity Regularizer not working with quantization aware training (QAT)

Describe the bug
Activity Regularizer not working with quantization aware training (QAT). On further evaluation I also saw that it creates a tensor of type <class 'tensorflow.python.framework.ops.Tensor'> for Actvity regularization loss where as for bias and kernel regularization loss it creates eager tensor which is easy to access.

System information

TensorFlow version (installed from source or binary): TF 2.3

TensorFlow Model Optimization version (installed from source or binary): pip install

Python version: 3.6.9

Describe the expected behavior
The Activity regularization loss is also supposed to be a eager tensor that has numpy values associated with it just like bias and kernel. This happens only on quantized models, non-quantized model have eager activity regularization loss.

Describe the current behavior
Currently the activity regularization loss tensor is of type <class 'tensorflow.python.framework.ops.Tensor'> which makes the values inaccessible even when the train function is running in eager mode.

Code to reproduce the issue

import numpy as np
import tensorflow as tf
from tensorflow_model_optimization.python.core.quantization.keras import quantize
from tensorflow.python import keras
l = keras.layers

tf.config.run_functions_eagerly(True)

def layers_list():
  return [
      l.Conv2D(32, 5, padding='same', activation='relu',
               input_shape=image_input_shape(), activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001)),
      l.MaxPooling2D((2, 2), (2, 2), padding='same'),
      # TODO(pulkitb): Add BatchNorm when transformations are ready.
      # l.BatchNormalization(),
      l.Conv2D(64, 5, padding='same', activation='relu', activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001)),
      l.MaxPooling2D((2, 2), (2, 2), padding='same'),
      l.Flatten(),
      l.Dense(1024, activation='relu'),
      l.Dropout(0.4),
      l.Dense(10, activation='softmax')
  ]


def sequential_model():
  return keras.Sequential(layers_list())


def functional_model():
  """Builds an MNIST functional model."""
  inp = keras.Input(image_input_shape())
  x = l.Conv2D(32, 5, padding='same', activation='relu', activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(inp)
  x = l.MaxPooling2D((2, 2), (2, 2), padding='same')(x)
  # TODO(pulkitb): Add BatchNorm when transformations are ready.
  # x = l.BatchNormalization()(x)
  x = l.Conv2D(64, 5, padding='same', activation='relu', activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(x)
  x = l.MaxPooling2D((2, 2), (2, 2), padding='same')(x)
  x = l.Flatten()(x)
  x = l.Dense(1024, activation='relu')(x)
  x = l.Dropout(0.4)(x)
  out = l.Dense(10, activation='softmax')(x)

  return keras.models.Model([inp], [out])


def image_input_shape(img_rows=28, img_cols=28):
  if tf.keras.backend.image_data_format() == 'channels_first':
    return 1, img_rows, img_cols
  else:
    return img_rows, img_cols, 1

def preprocessed_data(img_rows=28,
                      img_cols=28,
                      num_classes=10):
  """Get data for mnist training and evaluation."""
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

  if tf.keras.backend.image_data_format() == 'channels_first':
    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
  else:
    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)

  x_train = x_train.astype('float32')
  x_test = x_test.astype('float32')
  x_train /= 255
  x_test /= 255

  # convert class vectors to binary class matrices
  y_train = tf.keras.utils.to_categorical(y_train, num_classes)
  y_test = tf.keras.utils.to_categorical(y_test, num_classes)

  return x_train, y_train, x_test, y_test



model = functional_model() #sequential_model()
model.summary()
x_train, y_train, x_test, y_test = preprocessed_data()

model.compile(
    loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
model.fit(x_train, y_train, batch_size=500)
_, model_accuracy = model.evaluate(x_test, y_test, verbose=0)

print("Quantizing model")

quantized_model = quantize.quantize_model(model)
print(quantized_model.losses)
quantized_model.compile(
    loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
print(tf.math.add_n(model.losses).numpy())

quantized_model.fit(x_train, y_train, batch_size=500)
_, quantized_model_accuracy = quantized_model.evaluate(
    x_test, y_test, verbose=0)

Error Output

  1/120 [..............................] - ETA: 0s - loss: 2.3153 - accuracy: 0.1040WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0017s vs `on_train_batch_end` time: 0.0110s). Check your callbacks.
120/120 [==============================] - 1s 11ms/step - loss: 2.2161 - accuracy: 0.3372
Quantizing model
[<tf.Tensor 'conv2d/ActivityRegularizer_2/truediv:0' shape=() dtype=float32>, <tf.Tensor: shape=(), dtype=float32, numpy=0.00021372623>, <tf.Tensor 'conv2d_1/ActivityRegularizer_2/truediv:0' shape=() dtype=float32>, <tf.Tensor: shape=(), dtype=float32, numpy=0.004322933>]
Traceback (most recent call last):
  File "<path to python file>/tempTrain.py", line 95, in <module>
    print(quantized_model.losses.numpy())
AttributeError: 'list' object has no attribute 'numpy'

Additional context
I have reported this issue on tensorflow main github issues too as I am not sure which team is responsible for this.
Link to issue raised in tf github - [Tensorflow Issue] (https://github.com/tensorflow/tensorflow/issues/51680). Update, person assigned asked me to report it with model-optimization team. SO the above issue has been closed.

Start by running the provided reproduction with TensorFlow Model Optimization's quantize.quantize_model entry point and inspect the resulting model.losses values. Compare the activity regularizer losses with the kernel and bias regularizer losses under eager execution. Done means quantized models expose activity regularization losses as accessible eager tensors without the reported Tensor error.