Can't deserialize model using __add__ or Add()

Clone of https://github.com/tensorflow/tensorflow/issues/57574

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• TensorFlow installed from (source or binary): source
• TensorFlow version (use command below): v2.9.1-0-gd8ce9f9c301 2.9.1
• Python version: 3.9.13
• Bazel version (if compiling from source): 5.0.0
• GPU model and memory: nvidia A100-SXM4-40GB
• Exact command to reproduce: python repro.py

Describe the problem.

When I declare a model as such:

import tensorflow as tf
from tensorflow.keras.layers import Dense, Layer

class MyLayer(Layer):
    def __init__(self):
        super().__init__()
        self.dense = Dense(512)
    def __call__(self, X):
        short = X
        X = self.dense(X)
        X = short + X
        return X


def main():
    model = tf.keras.Sequential([MyLayer()])
    model.build([None, 512])
    model.save("/tmp/my_model")

    tf.keras.models.load_model("/tmp/my_model")

if __name__ == "__main__":
    main()

Describe the current behavior.

I get an exception trying to load the model:

api_dispatcher.Dispatch=<bound method PyCapsule.Dispatch of <Dispatch(_add_dispatch): >, args=(<tf.Tensor 'Placeholder:0' shape=(None, 512) dtype=float32>,), kwargs={}
Traceback (most recent call last):
  File "/mnt/disks/data/src/repro.py", line 24, in <module>
    main()
  File "/mnt/disks/data/src/repro.py", line 21, in main
    tf.keras.models.load_model("/tmp/my_model")
  File "/mnt/disks/data/src/conda/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/mnt/disks/data/src/conda/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py", line 1078, in op_dispatch_handler
    result = api_dispatcher.Dispatch(args, kwargs)
TypeError: Missing required positional argument

Describe the expected behavior.

The model loads fine.

Contributing.

• Do you want to contribute a PR? (yes/no): no
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue.

notebook

[ianstenbit]

Hi Froody!

It looks like you are overriding __call__ directly instead of the call method, which is what subclasses of Layer should override (see Layer docs)

I modified the gist from tlakrayal@ to rename your override from __call__ to call and it now appears to be working as intended.

I hope this unblocks you for now, and I am going to look into options for how to better prevent this for future users.

Thanks!

[ianstenbit]

I don’t think there is a clean way of preventing accidental override of __call__ because there are valid use cases for overriding it.

I’ve looked through the Layer docs again and between the docstrings about call and __call__ and the examples, I think there is currently sufficient guidance toward overriding call for use cases like this one.