How to apply tfmot to a non tf.keras.Model type model..

Hi,
This is more of a question than a feature request which I don't know where else to post.
So I'm trying to perform quantization aware training to a model that's not of tf.keras.Model type but of a wrapper Model class with some forward(), loss(), and trainable_variables() functions and some layers of tf.keras.Layer() type. The backpropagation is done with tf.GraientTape and optimizer.apply_gradients() outside this Model class.
Is this possible to do? I was thinking perhaps with some lower-level functions I can find in this repo this is doable but can't say for sure and don't really know how much work it would be. Would be great if I can get some advice directly from the developers.
Thank you very much.

Start by reviewing the TensorFlow Model Optimization Toolkit quantization-aware training APIs and their tf.keras.Model and tf.keras.Layer integration points. Compare those with the wrapper's forward(), loss(), trainable_variables(), GradientTape, and optimizer.apply_gradients() flow. Done would be a documented supported path or a clear statement that this model type is unsupported.