`TokenAndPositionEmbedding`, `TransformerEncoder` and `TransformerDecoder` can be saved, but prevent the model from being loaded

[DavidLandup0]

Describe the bug Some of the new custom layers correctly override the get_config() method, so as to be able to be saved, but when a model is saved in h5 format - they can't be loaded.

To Reproduce https://colab.research.google.com/drive/1vjyTkuXCV6FTpHcP-_f41BxwizxpYe1q?usp=sharing

Expected behavior The constituent layers not being required to be passed in as custom objects.

Additional context While this could be possibly worked around - it just doesn't look like intended behavior :) The colab uses TF and Keras versions 2.8.2 and 2.8.0 respectively, but this happens on my local machine with 2.9.1 and 2.9.0 as well.

Would you like to help us fix it? Sure!

[chenmoneygithub]

@DavidLandup0 Thank you for reporting the issue!

To clarify, the saving/loading works with Keras saved model format, but breaks with h5 format. I checked the error, and it implies TokenAndPositionEmbedding should be passed as a custom_object. Could you share your proposed fix? Thanks!

@mattdangerw h5 is the legacy saving format from TF1 IIUC, do we want to support it in kerasNLP? if so, we should enforce unit test on this format.

[mattdangerw]

Thanks for filing!

I think we could clear up this issue by adding keras.utils.register_keras_serializable(package="keras_nlp") annotations to our layers. This would support h5, but also force our tf-style saved model loading into the code paths that first restore layers from configs, and then restore weights. Overall, I think this probably makes sense, and is consistent with what keras_cv is doing currently.