[RNN] Stateful LSTM can't be converted to TF Lite with Integer Quantization

1. System information

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

2. Code

Please check my code below. Please put 'LSTMlayer.onnx' file into folder 'path' (or modify the path).
This code fails with the described issue, but if you will change QUANTIZATION = 'None', it will work.

import tensorflow as tf
import os
import numpy as np

path = "\test"
QUANTIZATION = 'IntegerWithFloatFallback' # 'IntegerWithFloatFallback' or 'None'

def representative_dataset():
        dummy = np.zeros((1,1,512), dtype = np.float32)
        yield [dummy]

converter = tf.lite.TFLiteConverter.from_saved_model(path)

# we need experimental_enable_resource_variables and "select TensorFlow ops" for
# AssignVariableOp, ReadVariableOp, VarHandleOp operations, otherwise you will get an error
converter.experimental_enable_resource_variables = True
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS,
  tf.lite.OpsSet.SELECT_TF_OPS
]

if (QUANTIZATION == 'IntegerWithFloatFallback'):
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = representative_dataset

tflite_model = converter.convert()
with open(os.path.join(path, 'model.tflite'), 'wb') as f:
  f.write(tflite_model)

3. Failure after conversion

File "...\git_issue_code.py", line 26, in
tflite_model = converter.convert()
File "...\lib\site-packages\tensorflow\lite\python\lite.py", line 763, in wrapper
return self._convert_and_export_metrics(convert_func, *args, **kwargs)
File "...\lib\site-packages\tensorflow\lite\python\lite.py", line 749, in _convert_and_export_metrics
result = convert_func(self, *args, **kwargs)
File "...\lib\site-packages\tensorflow\lite\python\lite.py", line 1031, in convert
return self._optimize_tflite_model(
File "...\lib\site-packages\tensorflow\lite\python\convert_phase.py", line 226, in wrapper
raise error from None # Re-throws the exception.
File "...\lib\site-packages\tensorflow\lite\python\convert_phase.py", line 216, in wrapper
return func(*args, **kwargs)
File "...\lib\site-packages\tensorflow\lite\python\lite.py", line 714, in _optimize_tflite_model
model = self._quantize(
File "...\lib\site-packages\tensorflow\lite\python\lite.py", line 517, in _quantize
calibrate_quantize = _calibrator.Calibrator(result,
File "...\lib\site-packages\tensorflow\lite\python\optimize\calibrator.py", line 78, in init
raise ValueError("Failed to parse the model: %s." % e)
ValueError: Failed to parse the model: Op FlexVarHandleOp missing inputs.

I have also reproduced this issue using the Keras LSTM layer directly. AssignVariableOp, ReadVariableOp, VarHandleOp are needed when you use stateful=True for LSTM layer (that is very important option for "infinite" data (e.g. audio stream)).
Please contact me if you need some additional info.

Thanks,
Best regards,
Anastasiia

Start at TFLiteConverter.convert and the quantization path leading to calibrator.py, using the provided stateful LSTM reproduction with IntegerWithFloatFallback. Investigate why calibration rejects FlexVarHandleOp with missing inputs; done means the same stateful LSTM model converts successfully with integer quantization while preserving the required resource-variable and select-TF-ops configuration.