How to save&load tf.keras.Model/Layer with extended methods/properties

Question

I want to extend tf.keras.Model to implement invertible neural networks, and save/load this model by tf.keras.Model.save() & tf.keras.mod