Add consistent support for name serialization of Operations. (#21373)

This is to solve an inconsistency in the saving / reloading of ops. Some ops behave differently.

Consider this code:

```python
input = keras.layers.Input(shape=(4,), dtype="float32")
output = keras.ops.abs(input)
model = keras.models.Model(input, output)

json = model.to_json()
reloaded_model = keras.models.model_from_json(json)
reloaded_json = reloaded_model.to_json()
```

The reloaded model JSON is the same as the original model JSON. The `abs` op is serialized as

- `{"module": "keras.src.ops.numpy", "class_name": "Absolute", "config": {"name": "absolute"}, "registered_name": "Absolute", "name": "absolute", ...}`

Consider the same code with `abs` replaced with `sum`:

```python
input = keras.layers.Input(shape=(4,), dtype="float32")
output = keras.ops.abs(input)
model = keras.models.Model(input, output)

json = model.to_json()
reloaded_model = keras.models.model_from_json(json)
reloaded_json = reloaded_model.to_json()
```

The reloaded model JSON is different from the original JSON.

- `{"module": "keras.src.ops.numpy", "class_name": "Sum", "config": {"axis": null, "keepdims": false}, "registered_name": "Sum", "name": "sum", ...}`
- `{"module": "keras.src.ops.numpy", "class_name": "Sum", "config": {"axis": null, "keepdims": false}, "registered_name": "Sum", "name": "sum_1", ...}`

The reloaded `sum` op now has a name `"sum_1"` instead of `"sum"`.

This is because:
- `Abs` does not define `__init__` and inherits the `Operation.__init__` that has a `name` parameter.
- `Sum` defines `__init__` without a `name` parameter.

We want saving / reloading to be idempotent. Even though users cannot control the name of ops (although it would be easy to add), the auto-assigned names should be saved and reloaded. For this, `name` has to be supported in `__init__`.

This PR changes the behavior of serialization from `_auto_config` to allow `name` to be returned when the `__init__` signature has `name` or `**kwargs`. The existing logic had a bug. Also switched to `inspect.signature` which is the recommended API in Python 3.

Note that adding `name` as a parameter to all ops `__init__`s will be done as a separate PR.

Author: hertschuh

keras/src/ops/operation.py

@@ -171,13 +171,16 @@ def get_config(self):
         # In this case the subclass doesn't implement get_config():
         # Let's see if we can autogenerate it.
         if getattr(self, "_auto_config", None) is not None:
-            xtra_args = set(config.keys())
             config.update(self._auto_config.config)
-            # Remove args non explicitly supported
-            argspec = inspect.getfullargspec(self.__init__)
-            if argspec.varkw != "kwargs":
-                for key in xtra_args - xtra_args.intersection(argspec.args[1:]):
-                    config.pop(key, None)
+            init_params = inspect.signature(self.__init__).parameters
+            init_has_name = "name" in init_params
+            init_has_kwargs = (
+                "kwargs" in init_params
+                and init_params["kwargs"].kind == inspect.Parameter.VAR_KEYWORD
+            )
+            if not init_has_name and not init_has_kwargs:
+                # We can't pass `name` back to `__init__`, remove it.
+                config.pop("name", None)
             return config
         else:
             raise NotImplementedError(

keras/src/ops/operation_test.py

@@ -30,20 +30,71 @@ def compute_output_spec(self, x):
 
 
 class OpWithCustomConstructor(operation.Operation):
-    def __init__(self, alpha, mode="foo"):
+    def __init__(self, alpha, *, name=None):
+        super().__init__(name=name)
+        self.alpha = alpha
+
+    def call(self, x):
+        return self.alpha * x
+
+    def compute_output_spec(self, x):
+        return keras_tensor.KerasTensor(x.shape, x.dtype)
+
+
+class OpWithCustomConstructorNoName(operation.Operation):
+    def __init__(self, alpha):
         super().__init__()
         self.alpha = alpha
-        self.mode = mode
 
     def call(self, x):
-        if self.mode == "foo":
-            return x
         return self.alpha * x
 
     def compute_output_spec(self, x):
         return keras_tensor.KerasTensor(x.shape, x.dtype)
 
 
+class OpWithKwargsInConstructor(operation.Operation):
+    def __init__(self, alpha, **kwargs):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+
+    def call(self, x):
+        return self.alpha * x
+
+    def compute_output_spec(self, x):
+        return keras_tensor.KerasTensor(x.shape, x.dtype)
+
+
+class OpWithCustomConstructorGetConfig(operation.Operation):
+    def __init__(self, alpha, *, name=None):
+        super().__init__(name=name)
+        self.alpha = alpha
+
+    def call(self, x):
+        return self.alpha * x
+
+    def compute_output_spec(self, x):
+        return keras_tensor.KerasTensor(x.shape, x.dtype)
+
+    def get_config(self):
+        return {**super().get_config(), "alpha": self.alpha}
+
+
+class OpWithKwargsInConstructorGetConfig(operation.Operation):
+    def __init__(self, alpha, **kwargs):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+
+    def call(self, x):
+        return self.alpha * x
+
+    def compute_output_spec(self, x):
+        return keras_tensor.KerasTensor(x.shape, x.dtype)
+
+    def get_config(self):
+        return {**super().get_config(), "alpha": self.alpha}
+
+
 class OperationTest(testing.TestCase):
     def test_symbolic_call(self):
         x = keras_tensor.KerasTensor(shape=(2, 3), name="x")
@@ -129,19 +180,100 @@ def test_eager_call(self):
         self.assertAllClose(out[0], np.ones((2, 3)))
         self.assertAllClose(out[1], np.ones((2, 3)) + 1)
 
-    def test_serialization(self):
-        op = OpWithMultipleOutputs(name="test_op")
+    def test_serialization_with_default_init_and_get_config(self):
+        # Explicit name passed in constructor is serialized and deserialized.
+        op = OpWithMultipleInputs(name="test_op")
         config = op.get_config()
         self.assertEqual(config, {"name": "test_op"})
-        op = OpWithMultipleOutputs.from_config(config)
-        self.assertEqual(op.name, "test_op")
+        revived = OpWithMultipleInputs.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+        # Auto generated name is serialized and deserialized.
+        op = OpWithMultipleInputs()
+        config = op.get_config()
+        self.assertEqual(config, {"name": op.name})
+        revived = OpWithMultipleInputs.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+    def test_serialization_custom_constructor_with_name_auto_config(self):
+        # Explicit name passed in constructor is serialized and deserialized.
+        op = OpWithCustomConstructor(alpha=0.2, name="test_op")
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": "test_op"})
+        revived = OpWithCustomConstructor.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
 
-    def test_autoconfig(self):
-        op = OpWithCustomConstructor(alpha=0.2, mode="bar")
+        # Auto generated name is serialized and deserialized.
+        op = OpWithCustomConstructor(alpha=0.2)
         config = op.get_config()
-        self.assertEqual(config, {"alpha": 0.2, "mode": "bar"})
+        self.assertEqual(config, {"alpha": 0.2, "name": op.name})
         revived = OpWithCustomConstructor.from_config(config)
         self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+    def test_serialization_custom_constructor_with_no_name_auto_config(self):
+        # Auto generated name is not serialized.
+        op = OpWithCustomConstructorNoName(alpha=0.2)
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2})
+        revived = OpWithCustomConstructorNoName.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+
+    def test_serialization_custom_constructor_with_kwargs_auto_config(self):
+        # Explicit name passed in constructor is serialized and deserialized.
+        op = OpWithKwargsInConstructor(alpha=0.2, name="test_op")
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": "test_op"})
+        revived = OpWithKwargsInConstructor.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+        # Auto generated name is serialized and deserialized.
+        op = OpWithKwargsInConstructor(alpha=0.2)
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": op.name})
+        revived = OpWithKwargsInConstructor.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+    def test_serialization_custom_constructor_custom_get_config(self):
+        # Explicit name passed in constructor is serialized and deserialized.
+        op = OpWithCustomConstructorGetConfig(alpha=0.2, name="test_op")
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": "test_op"})
+        revived = OpWithCustomConstructorGetConfig.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+        # Auto generated name is serialized and deserialized.
+        op = OpWithCustomConstructorGetConfig(alpha=0.2)
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": op.name})
+        revived = OpWithCustomConstructorGetConfig.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+    def test_serialization_custom_constructor_with_kwargs_custom_get_config(
+        self,
+    ):
+        # Explicit name passed in constructor is serialized and deserialized.
+        op = OpWithKwargsInConstructorGetConfig(alpha=0.2, name="test_op")
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": "test_op"})
+        revived = OpWithKwargsInConstructorGetConfig.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
+
+        # Auto generated name is serialized and deserialized.
+        op = OpWithKwargsInConstructorGetConfig(alpha=0.2)
+        config = op.get_config()
+        self.assertEqual(config, {"alpha": 0.2, "name": op.name})
+        revived = OpWithKwargsInConstructorGetConfig.from_config(config)
+        self.assertEqual(revived.get_config(), config)
+        self.assertEqual(revived.name, op.name)
 
     @skip_if_backend(
         "openvino", "Can not constant fold eltwise node by CPU plugin"