Keras TFOpLambda may conflict with quantization

Describe the bug

When we try to quantize a model containing a TFOpLambda layer, an AttributeError 'list' object has no attribute 'dtype' would occur.

System information

TensorFlow version (installed from binary): 2.5.0

TensorFlow Model Optimization version (installed from binary): 0.6.0

Python version: 3.9.6

Describe the expected behavior

Should return a quantized model without any error.

Describe the current behavior

AttributeError: 'list' object has no attribute 'dtype'

Code to reproduce the issue

https://colab.research.google.com/gist/hguandl/740df0dd42be2b220be26671c1e33119/oplambda-repro-minimal.ipynb

import tensorflow_model_optimization as tfmot
from tensorflow import keras
from tensorflow.python.keras.layers.core import TFOpLambda
from tensorflow.python.ops.math_ops import _add_dispatch
from tensorflow.python.util.tf_export import tf_export


@tf_export("test_lambda")
def custom_layer(tensor):
    return _add_dispatch(tensor, 2)


inputs = keras.Input(shape=(784,))
outputs = TFOpLambda(custom_layer)(inputs)

model = keras.models.Model(inputs=inputs, outputs=outputs)

q_model = tfmot.quantization.keras.quantize_model(model)

Screenshots

Please refer to results in the Colab above.

Additional context

The bug was first discussed in https://github.com/tensorflow/model-optimization/issues/546. I have found that TFOpLambda layers are forced to enable _preserve_input_structure_in_config which prevents the unwrapping of single-tensor lists. Therefore the input of lambda is a List rather than a Tensor type.

About TFOpLambda: tf_op_layer.py;

About unwrapping: funcional.py.