Untracked embedding layer inside Sequential keras 3 model

[markomitos]

During tracking of variables ther is an issue where an embedding layer in a keras 3 model is not tracked properly, whereas the same model in keras 2 is tracked properly.

Here is the error message:

AssertionError: Tried to export a function which references an 'untracked' resource. TensorFlow objects (e.g. tf.Variable) captured by functions must be 'tracked' by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly. See the information below:
	Function name = b'__inference_signature_wrapper_604'
	Captured Tensor = <ResourceHandle(name="sequential/embedding/embeddings/8", device="/job:localhost/replica:0/task:0/device:CPU:0", container="Anonymous", type="tensorflow::Var", dtype and shapes : "[ DType enum: 1, Shape: [10,5] ]")>
	Trackable referencing this tensor = <tf.Variable 'sequential/embedding/embeddings:0' shape=(10, 5) dtype=float32>
	Internal Tensor = Tensor("600:0", shape=(), dtype=resource)

This is model function used for keras 2 and keras 3 model respectably:

def build_embedding_keras_model(vocab_size=10):
  """Builds a test model with an embedding initialized to one-hot vectors."""
  keras_model = tf_keras.models.Sequential()
  keras_model.add(tf_keras.layers.Embedding(input_dim=vocab_size, output_dim=5,
                                            embeddings_initializer=keras.initializers.RandomUniform(seed=42)))
  keras_model.add(tf_keras.layers.Softmax())
  return keras_model

def build_embedding_keras3_model(vocab_size=10):
  """Builds a test model with an embedding initialized to one-hot vectors."""
  keras_model = keras.models.Sequential()
  keras_model.add(keras.layers.Embedding(input_dim=vocab_size, output_dim=5,
                                         embeddings_initializer=keras.initializers.RandomUniform(seed=42)))
  keras_model.add(keras.layers.Softmax())
  return keras_model

The tensor called "handle" inside embedding seems to be not tracked properly:

Here is the func graph after tracking:

The error is raised inside the ExportedConcreteFunction call during mapping of captured tensors:

https://github.com/tensorflow/tensorflow/blob/96a931bb3e145719ae111507f004b151a653027d/tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py#L45

[mehtamansi29]

Hi @markomitos -

Thanks for reporting the issue.
Here both the model has same weights, layers and generating same output as well. Attached gist here for the reference.

The error is raised inside the ExportedConcreteFunction call during mapping of captured tensors:
https://github.com/tensorflow/tensorflow/blob/96a931bb3e145719ae111507f004b151a653027d/tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py#L45

But the error you getting is regarding tensorflow repo. So can you please raise the new issue in tensorflow repo ?