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