SavedModel format doesn't support masking in Keras 3

[drasmuss]

Loading a SavedModel in Keras 3 loses any mask information on the model output. Loading a SavedModel in tf-keras preserves the mask information, as expected. Here is an example:

USE_KERAS_3 = True

import tensorflow as tf
if USE_KERAS_3:
    import keras
else:
    import tf_keras as keras



inp = keras.Input((1,))
x = keras.layers.Masking(mask_value=0.0)(inp)
model = keras.Model(inp, x)

if USE_KERAS_3:
    model.export("saved_model")
    loaded = keras.layers.TFSMLayer("saved_model")
else:
    model.save("saved_model", save_format="tf")
    loaded = keras.models.load_model("saved_model")


print(loaded(tf.zeros((1, 1)))._keras_mask)

With USE_KERAS_3=True this gives:

AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_keras_mask'

With USE_KERAS_3=False this gives:

tf.Tensor([False], shape=(1,), dtype=bool)