Keras model with LSTM quantization aware training

I would like to run quantization aware training with a keras model which has an LSTM layer. However, just the LSTM layer seems to not be supported. Alan Chiao seems to suggest here that it is possible to use the Quantize Config to allow this layer to be used. However, I can’t seem to make the right Quantize Config which allows for the use of the LSTM layer.

I attach below both the code for my model and the Quantize Config.

Could anyone confirm either

1. how to properly make a quantization aware model with an LSTM layer with keras
2. if LSTM quantization simply isn’t supported with keras

Thanks for the feedback!

from tensorflow.keras import layers as L
quant_layer = tfmot.quantization.keras.quantize_annotate_layer
quant_apply = tfmot.quantization.keras.quantize_apply
quantize_scope = tfmot.quantization.keras.quantize_scope

######################  Simple Quantized Model   ####################

i = tf.keras.Input(shape=(32,32,1))
x = (L.Conv2D(1, (5, 1), activation='relu', padding='same'))(i)
x = L.BatchNormalization()(x)
x = L.Lambda(lambda q: K.squeeze(q, -1), name='squeeze_last_dim')(x)
x = quant_layer(L.LSTM(64, return_sequences=True),DefaultLSTMQuantizeConfig())(x)
x = tf.keras.layers.Flatten()(x)
x = L.Dense(64, activation='relu')(x)
o = L.Dense(12, activation='softmax')(x)

simple_model = Model(inputs=[i], outputs=[o])
simple_model.compile(optimizer='adam', loss=['sparse_categorical_crossentropy'], metrics=['sparse_categorical_accuracy'])

with quantize_scope(
    {'DefaultLSTMQuantizeConfig': DefaultLSTMQuantizeConfig}):
    simple_model_quant = quant_apply(simple_model)
     
simple_model_quant.summary()

##################### Quantize Config  #################
LastValueQuantizer = tfmot.quantization.keras.quantizers.LastValueQuantizer
MovingAverageQuantizer = tfmot.quantization.keras.quantizers.MovingAverageQuantizer

class DefaultLSTMQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
    # Configure how to quantize weights.
    def get_weights_and_quantizers(self, layer):
        return [(layer.layer.recurrent_kernel, LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False))]

    # Configure how to quantize activations.
    def get_activations_and_quantizers(self, layer):
        return [(layer.activation, MovingAverageQuantizer(num_bits=8, symmetric=False, narrow_range=False, per_axis=False))]

    def set_quantize_weights(self, layer, quantize_weights):
      # Add this line for each item returned in `get_weights_and_quantizers`
      # , in the same order
        layer.kernel = quantize_weights[0]

    def set_quantize_activations(self, layer, quantize_activations):
      # Add this line for each item returned in `get_activations_and_quantizers`
      # , in the same order.
        layer.activation = quantize_activations[0]

    # Configure how to quantize outputs (may be equivalent to activations).
    def get_output_quantizers(self, layer):
        return []

    def get_config(self):
        return {}