Keras Cannot call legacy (saved with Keras 2) SavedModel models that have multiple inputs on Keras 3

When loading models saved in the legacy SavedModel format (TF/Keras 2) in a Keras 3 environment, I can successfully create a TFSMLayer, but I'm unable to call it when the original model had more than one input (an array of 2 inputs, in my case). No matter the combination of input arguments, I get errors that resemble this:

>>> # python 3.11, tensorflow 2.18.0, keras 3.8.0 environment
>>> image_input = tf.random.normal((1,10))
>>> label_input = tf.random.normal((1,2))
>>>
>>> # loading the model that had been saved in keras2
>>> tf.keras.layers.TFSMLayer('temp/multi_input_keras2', call_endpoint='serving_default')(inputs=[image_input, label_input])

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "python3.11/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "python3.11/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py", line 1184, in _call_impl
    raise TypeError(  # pylint: disable=raise-missing-from
TypeError: Exception encountered when calling TFSMLayer.call().

Binding inputs to tf.function failed due to `too many positional arguments`. Received args: ([<tf.Tensor: shape=(1, 10), dtype=float32, numpy=
array([[0.99221754, 0.22724748, 0.46073282, 0.47727287, 0.13898277,
        0.8761684 , 0.86543703, 0.3392017 , 0.86744404, 0.40240443]],
      dtype=float32)>, <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[0.33422947, 0.21715128]], dtype=float32)>],) and kwargs: {} for signature: (*, image_input: TensorSpec
(shape=(None, 10), dtype=tf.float32, name='image_input'), label_input: TensorSpec(shape=(None, 2), dtype=tf.float32, name='label_input')) -> Dict[['dense_1', TensorSpec(shape=(None, 1), dtype=tf.float32, name='dense_1')]].
Fallback to flat signature also failed due to: signature_wrapper(image_input, label_input) takes 0 positional arguments, got 1.

Arguments received by TFSMLayer.call():
  • inputs=['tf.Tensor(shape=(1, 10), dtype=float32)', 'tf.Tensor(shape=(1, 2), dtype=float32)']
  • training=False
  • kwargs=<class 'inspect._empty'>
>>>

I have tried calling the model (loaded through legacy_model = TFSMLayer(path_to_keras2_saved_model, call_endpoint='serving_default')) in the following ways, without success: - legacy_model([image_input, label_input]) - legacy_model((image_input, label_input)) - legacy_model({"image_input": image_input, "label_input": label_input}) - legacy_model(inputs=[image_input, label_input]) - legacy_model(inputs=(image_input, label_input)) - legacy_model(inputs={"image_input": image_input, "label_input": label_input}) - All of these have a very similar message to the output above - legacy_model() - TypeError: missing a required argument: 'inputs'

If I load and call a model with the same signature (2 tensor inputs) but saved in Keras 3, it works as expected:

>>> image_input = tf.random.normal((1,10))
>>> label_input = tf.random.normal((1,2))
>>>
>>> # loading the model that had been saved in keras3
>>> tf.keras.layers.TFSMLayer('temp/multi_input_keras3')([image_input, label_input])
<tf.Tensor: shape=(1, 1), dtype=float32, numpy=array([[0.4987833]], dtype=float32)>
>>>

If I load a model saved with Keras 2 but that has only a single input, it works fine as well:

>>> image_input = tf.random.normal((1,10))
>>>
>>> # loading the model that had been saved in keras2
>>> tf.keras.layers.TFSMLayer('temp/single_input_keras2', call_endpoint='serving_default')(image_input)
{'dense_1': <tf.Tensor: shape=(1, 1), dtype=float32, numpy=array([[0.24625102]], dtype=float32)>}
>>>

To reproduce

The models used here have been attached (40kb), and the code I used is this:

import tensorflow as tf

# executed in tensorflow 2.10.1 (keras 2)
def save_model_in_keras2():
    input_layer_0 = tf.keras.layers.Input(shape=(10,), name="image_input")
    hidden_layer_1 = tf.keras.layers.Dense(10, activation='relu')(input_layer_0)
    output_layer = tf.keras.layers.Dense(1, activation='sigmoid')(hidden_layer_1)
    single_input_model = tf.keras.models.Model(inputs=input_layer_0, outputs=output_layer)
    single_input_model.save('temp/single_input_keras2')

    input_layer_1 = tf.keras.layers.Input(shape=(2,), name="label_input")
    hidden_layer_1 = tf.keras.layers.Concatenate()([input_layer_0, input_layer_1])
    hidden_layer_2 = tf.keras.layers.Dense(2, activation='relu')(hidden_layer_1)
    output_layer = tf.keras.layers.Dense(1, activation='sigmoid')(hidden_layer_2)
    multi_input_model = tf.keras.models.Model(inputs=(input_layer_0, input_layer_1), outputs=output_layer)
    multi_input_model.save('temp/multi_input_keras2')

# executed in tensorflow 2.18.0 (keras 3)
def save_model_in_keras3():
    input_layer_0 = tf.keras.layers.Input(shape=(10,), name="image_input")
    hidden_layer_1 = tf.keras.layers.Dense(10, activation='relu')(input_layer_0)
    output_layer = tf.keras.layers.Dense(1, activation='sigmoid')(hidden_layer_1)
    single_input_model = tf.keras.models.Model(inputs=input_layer_0, outputs=output_layer)
    single_input_model.export('temp/single_input_keras3')

    input_layer_1 = tf.keras.layers.Input(shape=(2,), name="label_input")
    hidden_layer_1 = tf.keras.layers.Concatenate()([input_layer_0, input_layer_1])
    hidden_layer_2 = tf.keras.layers.Dense(2, activation='relu')(hidden_layer_1)
    output_layer = tf.keras.layers.Dense(1, activation='sigmoid')(hidden_layer_2)
    multi_input_model = tf.keras.models.Model(inputs=(input_layer_0, input_layer_1), outputs=output_layer)
    multi_input_model.export('temp/multi_input_keras3')

To reproduce the error after unzipping the folder with the saved models:

# execute in tensorflow 2.18.0 (keras 3)
image_input = tf.random.normal((1, 10))
label_input = tf.random.normal((1, 2))

# load keras3 model in keras3 (works)
model = tf.keras.layers.TFSMLayer('temp/multi_input_keras3')
result = model([image_input, label_input])
print("result:", result)

# load legacy model in keras 3 (calling it fails)
legacy_model = tf.keras.layers.TFSMLayer('temp/multi_input_keras2', call_endpoint='serving_default')
result = legacy_model([image_input, label_input])
print("result:", result)

My environment

OS: Windows 10 running WSL2 on Ubuntu 24.04
Python 3.11
Tensorflow 2.18.0 (Keras 3.8.0)

Comment From: dhantule

Hi @fegemo, thanks for reporting this.

Calling a TSFM layer works with single input tensor, if your model requires multiple input tensors you should try subclassing the TSFM layer and implementing a call() method with a custom signature as mentioned here.