Add name parameter to the __init__ of all op classes. (#21376)

Follow up of #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 adds a `name` parameter to all existing `__init__` of op classes, which is passed to `super().__init__()`. Note that it is defined as a keyword only argument for forward compatibility in case more parameters get added in the feature.

Empty `__init__`s were removed.

Fix `UnravelIndex.__init__` which was not calling `super().__init__()`.

Author: hertschuh

keras/src/ops/core.py

@@ -12,9 +12,6 @@
 
 
 class Map(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, f, xs):
         return backend.core.map(f, xs)
 
@@ -78,8 +75,8 @@ def map(f, xs):
 
 
 class Scan(Operation):
-    def __init__(self, length=None, reverse=False, unroll=1):
-        super().__init__()
+    def __init__(self, length=None, reverse=False, unroll=1, *, name=None):
+        super().__init__(name=name)
         self.length = length
         self.reverse = reverse
         self.unroll = unroll
@@ -191,8 +188,8 @@ def scan(f, init, xs, length=None):
 
 
 class AssociativeScan(Operation):
-    def __init__(self, reverse=False, axis=0):
-        super().__init__()
+    def __init__(self, reverse=False, axis=0, *, name=None):
+        super().__init__(name=name)
         self.reverse = reverse
         self.axis = axis
 
@@ -289,8 +286,8 @@ def associative_scan(f, elems, reverse=False, axis=0):
 
 
 class Scatter(Operation):
-    def __init__(self, shape):
-        super().__init__()
+    def __init__(self, shape, *, name=None):
+        super().__init__(name=name)
         self.shape = shape
 
     def call(self, indices, values):
@@ -392,8 +389,8 @@ def scatter_update(inputs, indices, updates):
 
 
 class Slice(Operation):
-    def __init__(self, shape):
-        super().__init__()
+    def __init__(self, shape, *, name=None):
+        super().__init__(name=name)
         self.shape = shape
 
     def call(self, inputs, start_indices):
@@ -530,8 +527,8 @@ def switch(index, branches, *operands):
 
 
 class WhileLoop(Operation):
-    def __init__(self, cond, body, maximum_iterations=None):
-        super().__init__()
+    def __init__(self, cond, body, maximum_iterations=None, *, name=None):
+        super().__init__(name=name)
         self.cond = cond
         self.body = body
         self.maximum_iterations = maximum_iterations
@@ -599,9 +596,6 @@ def while_loop(
 
 
 class StopGradient(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, variable):
         return backend.core.stop_gradient(variable)
 
@@ -634,8 +628,8 @@ def stop_gradient(variable):
 
 
 class ForiLoop(Operation):
-    def __init__(self, lower, upper, body_fun):
-        super().__init__()
+    def __init__(self, lower, upper, body_fun, *, name=None):
+        super().__init__(name=name)
         self.lower = lower
         self.upper = upper
         self.body_fun = body_fun
@@ -682,8 +676,8 @@ def fori_loop(lower, upper, body_fun, init_val):
 
 
 class Unstack(Operation):
-    def __init__(self, num=None, axis=0):
-        super().__init__()
+    def __init__(self, num=None, axis=0, *, name=None):
+        super().__init__(name=name)
         self.num = num
         self.axis = axis
 
@@ -787,8 +781,8 @@ def dtype(x):
 
 
 class Cast(Operation):
-    def __init__(self, dtype):
-        super().__init__()
+    def __init__(self, dtype, *, name=None):
+        super().__init__(name=name)
         self.dtype = backend.standardize_dtype(dtype)
 
     def call(self, x):
@@ -822,8 +816,8 @@ def cast(x, dtype):
 
 
 class SaturateCast(Operation):
-    def __init__(self, dtype):
-        super().__init__()
+    def __init__(self, dtype, *, name=None):
+        super().__init__(name=name)
         self.dtype = backend.standardize_dtype(dtype)
 
     def call(self, x):
@@ -928,8 +922,8 @@ def get_dtype_min_max(dtype):
 
 
 class ConvertToTensor(Operation):
-    def __init__(self, dtype=None, sparse=None, ragged=None):
-        super().__init__()
+    def __init__(self, dtype=None, sparse=None, ragged=None, *, name=None):
+        super().__init__(name=name)
         self.dtype = None if dtype is None else backend.standardize_dtype(dtype)
         self.sparse = sparse
         self.ragged = ragged

keras/src/ops/image.py

@@ -8,8 +8,8 @@
 
 
 class RGBToGrayscale(Operation):
-    def __init__(self, data_format=None):
-        super().__init__()
+    def __init__(self, data_format=None, *, name=None):
+        super().__init__(name=name)
         self.data_format = backend.standardize_data_format(data_format)
 
     def call(self, images):
@@ -77,8 +77,8 @@ def rgb_to_grayscale(images, data_format=None):
 
 
 class RGBToHSV(Operation):
-    def __init__(self, data_format=None):
-        super().__init__()
+    def __init__(self, data_format=None, *, name=None):
+        super().__init__(name=name)
         self.data_format = backend.standardize_data_format(data_format)
 
     def call(self, images):
@@ -149,8 +149,8 @@ def rgb_to_hsv(images, data_format=None):
 
 
 class HSVToRGB(Operation):
-    def __init__(self, data_format=None):
-        super().__init__()
+    def __init__(self, data_format=None, *, name=None):
+        super().__init__(name=name)
         self.data_format = backend.standardize_data_format(data_format)
 
     def call(self, images):
@@ -228,8 +228,10 @@ def __init__(
         fill_mode="constant",
         fill_value=0.0,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.size = tuple(size)
         self.interpolation = interpolation
         self.antialias = antialias
@@ -413,8 +415,10 @@ def __init__(
         fill_mode="constant",
         fill_value=0,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.interpolation = interpolation
         self.fill_mode = fill_mode
         self.fill_value = fill_value
@@ -554,8 +558,10 @@ def __init__(
         dilation_rate=1,
         padding="valid",
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         if isinstance(size, int):
             size = (size, size)
         self.size = size
@@ -707,8 +713,8 @@ def _extract_patches(
 
 
 class MapCoordinates(Operation):
-    def __init__(self, order, fill_mode="constant", fill_value=0):
-        super().__init__()
+    def __init__(self, order, fill_mode="constant", fill_value=0, *, name=None):
+        super().__init__(name=name)
         self.order = order
         self.fill_mode = fill_mode
         self.fill_value = fill_value
@@ -803,8 +809,10 @@ def __init__(
         target_height=None,
         target_width=None,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.top_padding = top_padding
         self.left_padding = left_padding
         self.bottom_padding = bottom_padding
@@ -1014,8 +1022,10 @@ def __init__(
         target_height=None,
         target_width=None,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.top_cropping = top_cropping
         self.bottom_cropping = bottom_cropping
         self.left_cropping = left_cropping
@@ -1238,8 +1248,10 @@ def __init__(
         interpolation="bilinear",
         fill_value=0,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.interpolation = interpolation
         self.fill_value = fill_value
         self.data_format = backend.standardize_data_format(data_format)
@@ -1381,8 +1393,10 @@ def __init__(
         kernel_size=(3, 3),
         sigma=(1.0, 1.0),
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.kernel_size = kernel_size
         self.sigma = sigma
         self.data_format = backend.standardize_data_format(data_format)
@@ -1470,8 +1484,10 @@ def __init__(
         fill_value=0.0,
         seed=None,
         data_format=None,
+        *,
+        name=None,
     ):
-        super().__init__()
+        super().__init__(name=name)
         self.alpha = alpha
         self.sigma = sigma
         self.interpolation = interpolation

keras/src/ops/linalg.py

@@ -7,9 +7,6 @@
 
 
 class Cholesky(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _cholesky(x)
 
@@ -47,9 +44,6 @@ def _cholesky(x):
 
 
 class Det(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _det(x)
 
@@ -83,9 +77,6 @@ def _det(x):
 
 
 class Eig(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _eig(x)
 
@@ -122,9 +113,6 @@ def _eig(x):
 
 
 class Eigh(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _eigh(x)
 
@@ -162,9 +150,6 @@ def _eigh(x):
 
 
 class Inv(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _inv(x)
 
@@ -198,9 +183,6 @@ def _inv(x):
 
 
 class LuFactor(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, x):
         return _lu_factor(x)
 
@@ -248,8 +230,8 @@ def _lu_factor(x):
 
 
 class Norm(Operation):
-    def __init__(self, ord=None, axis=None, keepdims=False):
-        super().__init__()
+    def __init__(self, ord=None, axis=None, keepdims=False, *, name=None):
+        super().__init__(name=name)
         if isinstance(ord, str):
             if ord not in ("fro", "nuc"):
                 raise ValueError(
@@ -367,8 +349,8 @@ def norm(x, ord=None, axis=None, keepdims=False):
 
 
 class Qr(Operation):
-    def __init__(self, mode="reduced"):
-        super().__init__()
+    def __init__(self, mode="reduced", *, name=None):
+        super().__init__(name=name)
         if mode not in {"reduced", "complete"}:
             raise ValueError(
                 "`mode` argument value not supported. "
@@ -440,9 +422,6 @@ def qr(x, mode="reduced"):
 
 
 class Solve(Operation):
-    def __init__(self):
-        super().__init__()
-
     def call(self, a, b):
         return _solve(a, b)
 
@@ -484,8 +463,8 @@ def _solve(a, b):
 
 
 class SolveTriangular(Operation):
-    def __init__(self, lower=False):
-        super().__init__()
+    def __init__(self, lower=False, *, name=None):
+        super().__init__(name=name)
         self.lower = lower
 
     def call(self, a, b):
@@ -531,8 +510,8 @@ def _solve_triangular(a, b, lower=False):
 
 
 class SVD(Operation):
-    def __init__(self, full_matrices=True, compute_uv=True):
-        super().__init__()
+    def __init__(self, full_matrices=True, compute_uv=True, *, name=None):
+        super().__init__(name=name)
         self.full_matrices = full_matrices
         self.compute_uv = compute_uv
 
@@ -586,8 +565,8 @@ def _svd(x, full_matrices=True, compute_uv=True):
 
 
 class Lstsq(Operation):
-    def __init__(self, rcond=None):
-        super().__init__()
+    def __init__(self, rcond=None, *, name=None):
+        super().__init__(name=name)
         self.rcond = rcond
 
     def call(self, a, b):