tf.quantization.quantize fails when converting to TFLite

1. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
• TensorFlow installation (pip package or built from source): from pip package tf-nightly
• TensorFlow library (version, if pip package or github SHA, if built from source): 2.7.0-dev20210922

2. Code

I'm exporting a TFLite model with multiple signatures provided as concrete functions. In one of them I want to perform a manual quantization using tf.quantization.quantize, since as far as I know the quantize operation exists in TFLite.

import tensorflow as tf

class TestModel(tf.keras.models.Model):
  @tf.function
  def quantize(self, value):
    # Range values are just an example for repro purposes.
    return tf.quantization.quantize(value, -1.0, 1.0, tf.qint8)


test_model = TestModel()
test_model.quantize(tf.random.uniform([10], -1.0, 1.0)) # Works fine.
signatures = [test_model.quantize.get_concrete_function(tf.TensorSpec([None, 10], tf.float32))]

converter = tf.lite.TFLiteConverter.from_concrete_functions(signatures, test_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

However, convert fails with the following error.

error: Failed to convert element type '!tf_type.qint8': Unsupported type
<unknown>:0: note: loc("StatefulPartitionedCall_2"): called from
<unknown>:0: error: invalid TFLite type: 'tensor<?x?x32x!tf_type.qint8>'

If instead I try to use tf.int8, then I get this other error.

TypeError: Value passed to parameter 'T' has DataType int8 not in list of allowed values: qint8, quint8, qint32, qint16, quint16

Since the operation actually exists in TFLite, could this be just a problem managing the output dtype argument?

It should be noted that I am fully aware of the inference_input_type and inference_output_type attributes in the converter. These are not what I'm asking about. I'm asking about explicitly running the quantize op within one of the model signatures.

Reproduce the failure using the supplied TestModel, tf.quantization.quantize, and TFLiteConverter.from_concrete_functions entry points, then inspect how converter.convert handles the qint8 output type. Confirm whether an explicit quantize operation can be converted in a model signature without the unsupported-type error, while preserving the requested output dtype.