Error deserialising model when Embedding layer size is defined using a tensor

[maxfisher-g]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Debian Unstable
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.15.0 (git version: unknown)
• Python version: 3.10.2
• GPU model and memory: N/A (reproducible on tensorflow CPU-only build)
• Exact command to reproduce: see code below

Describe the problem.

It's possible to create a model (specifically, an Embedding layer) with input_dim defined by a tf.constant, train it, and save the model to disk using Model.save(). However, deserialising it throws an exception.

Describe the current behavior.

The model is created and can be trained and serialised to disk without any errors occurring. However, deserialising throws an exception meaning that the model cannot be reloaded after serialisation.

Describe the expected behavior.

Either the model saves and loads without error, or an error is thrown during model construction to indicate that a Tensor cannot be used to define input_dim (or any other quantity for which a similar deserialisation issue would occur).

The expected saving/loading behaviour can be demonstrated by setting reproduce_bug = False or simply defining embedding_size = 16000 in the code below.

Contributing.

• Do you want to contribute a PR? (yes/no): No
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue.

import tensorflow as tf
from tensorflow.keras import layers

reproduce_bug = True
embedding_size = tf.constant(16000) if reproduce_bug else 16000

# Define model with embedding_size as a tf.constant
model = tf.keras.Sequential([
    layers.Input(shape=(512,), dtype=tf.int32),
    layers.Embedding(embedding_size, 128),
    layers.Bidirectional(layers.LSTM(128)),
    layers.Dense(1, activation='sigmoid'),   
])

# Model is created without error
model.summary()

# Model trains without error (code omitted)

# Model saves without error
model.save("test.keras", overwrite=True, save_format='keras')

# Loading model throws exception
tf.keras.models.load_model("test.keras").summary()

Source code / logs.

Traceback (most recent call last)
<ipython-input-13-2143c6fbea09> in <module>()
     19 model.save(save_filename, save_format='keras')
     20 
---> 21 tf.keras.models.load_model(save_filename).summary()

keras/saving/saving_api.py in load_model(filepath, custom_objects, compile, safe_mode, **kwargs)
    252                 f"with the native Keras format: {list(kwargs.keys())}"
    253             )
--> 254         return saving_lib.load_model(
    255             filepath,
    256             custom_objects=custom_objects,

keras/saving/saving_lib.py in load_model(filepath, custom_objects, compile, safe_mode)
    279 
    280     except Exception as e:
--> 281         raise e
    282     else:
    283         return model

keras/saving/saving_lib.py in load_model(filepath, custom_objects, compile, safe_mode)
    244             # Construct the model from the configuration file in the archive.
    245             with ObjectSharingScope():
--> 246                 model = deserialize_keras_object(
    247                     config_dict, custom_objects, safe_mode=safe_mode
    248                 )

keras/saving/serialization_lib.py in deserialize_keras_object(config, custom_objects, safe_mode, **kwargs)
    726     safe_mode_scope = SafeModeScope(safe_mode)
    727     with custom_obj_scope, safe_mode_scope:
--> 728         instance = cls.from_config(inner_config)
    729         build_config = config.get("build_config", None)
    730         if build_config:

keras/engine/sequential.py in from_config(cls, config, custom_objects)
    464         for layer_config in layer_configs:
    465             use_legacy_format = "module" not in layer_config
--> 466             layer = layer_module.deserialize(
    467                 layer_config,
    468                 custom_objects=custom_objects,

keras/layers/serialization.py in deserialize(config, custom_objects, use_legacy_format)
    274         )
    275 
--> 276     return serialization_lib.deserialize_keras_object(
    277         config,
    278         module_objects=LOCAL.ALL_OBJECTS,

keras/saving/serialization_lib.py in deserialize_keras_object(config, custom_objects, safe_mode, **kwargs)
    607                     custom_objects=custom_objects,
    608                 )
--> 609             return deserialize_keras_object(
    610                 serialize_with_public_class(
    611                     module_objects[config], inner_config=inner_config

keras/saving/serialization_lib.py in deserialize_keras_object(config, custom_objects, safe_mode, **kwargs)
    726     safe_mode_scope = SafeModeScope(safe_mode)
    727     with custom_obj_scope, safe_mode_scope:
--> 728         instance = cls.from_config(inner_config)
    729         build_config = config.get("build_config", None)
    730         if build_config:

keras/engine/base_layer.py in from_config(cls, config)
    868             return cls(**config)
    869         except Exception as e:
--> 870             raise TypeError(
    871                 f"Error when deserializing class '{cls.__name__}' using "
    872                 f"config={config}.\n\nException encountered: {e}"

TypeError: Error when deserializing class 'Embedding' using config={'name': 'embedding_9', 'trainable': True, 'dtype': 'float32', 'batch_input_shape': [None, None], 'input_dim': {'class_name': '__tensor__', 'config': {'value': 16000, 'dtype': 'int32'}}, 'output_dim': 128, 'embeddings_initializer': {'module': 'keras.initializers', 'class_name': 'RandomUniform', 'config': {'minval': -0.05, 'maxval': 0.05, 'seed': None}, 'registered_name': None}, 'embeddings_regularizer': None, 'activity_regularizer': None, 'embeddings_constraint': None, 'mask_zero': False, 'input_length': None}.

Exception encountered: '<=' not supported between instances of 'dict' and 'int'

[tilakrayal]

@sachinprasadhs,
I was able to reproduce the on tensorflow v2.12, v2.13 and tf-nightly. Kindly find the gist of it here.

[haifeng-jin]

Temporarily assigned to @nkovela1.
Please reassign if it is not relevant.