load_model no longer works with Dot layer in TensorFlow 2.16

[Dobiasd]

With TensorFlow 2.15.1, the following code works fine:

from tensorflow.keras.layers import Input, Dot
from tensorflow.keras.models import Model, load_model

input_1 = Input(shape=(2, 3, 4))
input_2 = Input(shape=(2, 3, 4))
x = Dot(axes=(1, 1))([input_1, input_2])
model = Model(inputs=[input_1, input_2], outputs=[x])
model.compile()
model.save('model.keras')

load_model('model.keras')

However, with TensorFlow 2.16.1, load_model throws an exception:

ValueError: A `Dot` layer should be called on a list of 2 inputs. Received: input_shape=[[None, 2, 3, 4], [None, 2, 3, 4]]

[tilakrayal]

@Dobiasd,
Yes, As you mentioned it was executed with the error on tensorflow v2.16 & keras 3.1.1. On tensorflow v2.15 it was executed without any error. And now it got fixed on tf-nightly(2.17.0-dev20240403) and keras(3.2.0.dev2024040303) which will be available in the upcoming releases.
Kindly find the attached gist for the reference.

Also there is PR raised for the similar issue in Keras repo which was merged recently.
keras-team/keras#19362

Thank you!

[Dobiasd]

Ok, perfect, thank you. 👍

[tilakrayal]

@Dobiasd,
As this issue is resolved in tf-nightly, Could you please feel free to move this issue to closed status. Thank you!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No