Incompatible shapes when loading model with TextVectorization and Embedding + Conv1D

[swierchola]

There is ValueError: Shapes incompatible when trying to load_model with TextVectorization and Embedding + Conv1D (or LSTM or GRU).

It can be easily replicated using TextVectorization layer examples from TF documentation. You just need to add model save & load to Using the TextVectorization layer in an Embedding + Conv1D model sample.
I created copy of TextVectorization layer examples colab with these errors. Code added to original example:

model.save('my_model')
loaded_model = tf.keras.models.load_model('my_model')

Error from load_model execution:

WARNING:tensorflow:Inconsistent references when loading the checkpoint into this object graph. 
Either the Trackable object references in the Python program have changed in an incompatible way, or the checkpoint was generated in an incompatible program.

Two checkpoint references resolved to different objects (<tensorflow.python.keras.layers.embeddings.Embedding object at 0x7f31fe4d8320> and <tensorflow.python.keras.saving.saved_model.load.TextVectorization object at 0x7f31fe4d8b70>).

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/tensor_shape.py in assert_is_compatible_with(self, other)
   1115     """
   1116     if not self.is_compatible_with(other):
-> 1117       raise ValueError("Shapes %s and %s are incompatible" % (self, other))
   1118 
   1119   def most_specific_compatible_shape(self, other):

ValueError: Shapes (128, 128) and (7, 128, 128) are incompatible

This error do NOT occur when trying to do save & load in example: Using the TextVectorization layer in a bigram TF-IDF densely-connected model. This example contains only Dense and Dropout layers and do not contain Embedding and Conv1D.

I also found same problem for my custom models using LSTM and GRU.

Problem occurred for nightly TF version 2.2.0-dev20200405

[gadagashwini-zz]

Was able to reproduce the issue with Tf == 2.2.0-dev20200406,
Please find the gist here. Thanks!

[imartinezl]

Was able to reproduce the issue with Tf == 2.2.0-dev20200406,
Please find the gist here. Thanks!

Hi @swierchola

I have loaded the model commented by @gadagashwini using tf-nightly==2.3.0-dev20200611, and it works without prompting any errors. The inference also works well. It would be great to know if you get the same results so that this issue can be closed.

[swierchola]

@imartinezl I confirm that problem is solved in latest tf-nightly. Probably commit Fix dimensionality handling issues in TextVectorization fixed it. Thanks.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[k-w-w]

In the latest nightly, please make sure that TextVectorization and the custom standardization function are registered when loading.

loaded_model = tf.keras.models.load_model(
    'my_model', 
    custom_objects={'TextVectorization':TextVectorization, 
                    'custom_standardization':custom_standardization})