Make ops consistent and test ops consistency. (#21334)

Added tests to verify that:
- All the parameters in the op class `call` (dynamic / tensor parameters) and `__init__` (static / non-tensor parameters) match the parameters in the op function
  - the names should match
  - default values should match
  - the order should be consistent, usually with the static parameters all at the end
- Ops are exported in `keras.ops`
- Ops are exported in their module, e.g. `keras.ops.numpy`
- Ops are implemented in every backend
- The backend implementation has the same signature as the op wrapper (same parameters in the same order with the same default values)

Note about default values: only the default values declared in the common op function are used by users, which is why those values were not changed and were applied to the class and backend implementations.
- The default values for the class `__init__` are not used because the op function passes all parameters when constructing the class.
- The default values for the backend implementation are not used, because the common op function passes all parameters when calling the backend implementation. The exception to this is when an op implementation calls another op directly.

Therefore:
- The defaults on the backend implementations are just a reminder for the implementer
- The defaults on the class `__init__` are mostly a reminder, but can also serve as a forward compatibility mechanism when deserializing a model after a parameter is added.

Fixed all inconsistencies, mostly inconsistent default values.

Author: hertschuh

keras/api/_tf_keras/keras/ops/numpy/__init__.py

@@ -133,6 +133,7 @@
 from keras.src.ops.numpy import roll as roll
 from keras.src.ops.numpy import rot90 as rot90
 from keras.src.ops.numpy import round as round
+from keras.src.ops.numpy import searchsorted as searchsorted
 from keras.src.ops.numpy import select as select
 from keras.src.ops.numpy import sign as sign
 from keras.src.ops.numpy import signbit as signbit

keras/api/ops/numpy/__init__.py

@@ -133,6 +133,7 @@
 from keras.src.ops.numpy import roll as roll
 from keras.src.ops.numpy import rot90 as rot90
 from keras.src.ops.numpy import round as round
+from keras.src.ops.numpy import searchsorted as searchsorted
 from keras.src.ops.numpy import select as select
 from keras.src.ops.numpy import sign as sign
 from keras.src.ops.numpy import signbit as signbit

keras/src/backend/jax/nn.py

@@ -155,9 +155,9 @@ def log_softmax(x, axis=-1):
     return jnn.log_softmax(x, axis=axis)
 
 
-def sparsemax(logits, axis=-1):
+def sparsemax(x, axis=-1):
     # Sort logits along the specified axis in descending order
-    logits = convert_to_tensor(logits)
+    logits = convert_to_tensor(x)
     logits_sorted = -1.0 * jnp.sort(logits * -1.0, axis=axis)
     logits_cumsum = jnp.cumsum(logits_sorted, axis=axis)  # find cumulative sum
     r = jnp.arange(1, logits.shape[axis] + 1)  # Determine the sparsity
@@ -250,8 +250,8 @@ def max_pool(
 def average_pool(
     inputs,
     pool_size,
-    strides,
-    padding,
+    strides=None,
+    padding="valid",
     data_format=None,
 ):
     data_format = backend.standardize_data_format(data_format)
@@ -485,7 +485,7 @@ def conv_transpose(
     )
 
 
-def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
+def one_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
     x = convert_to_tensor(x)
     if sparse:
         if axis < 0:
@@ -513,7 +513,7 @@ def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
     return jnn.one_hot(x, num_classes, axis=axis, dtype=dtype)
 
 
-def multi_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
+def multi_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
     x = convert_to_tensor(x)
     reduction_axis = 1 if len(x.shape) > 1 else 0
     if sparse:

keras/src/backend/jax/numpy.py

@@ -630,10 +630,10 @@ def digitize(x, bins):
     return jnp.digitize(x, bins)
 
 
-def dot(x, y):
-    x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
-    return jnp.dot(x, y)
+def dot(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.dot(x1, x2)
 
 
 def empty(shape, dtype=None):
@@ -982,9 +982,9 @@ def ravel(x):
     return jnp.ravel(x)
 
 
-def unravel_index(x, shape):
-    x = convert_to_tensor(x)
-    return jnp.unravel_index(x, shape)
+def unravel_index(indices, shape):
+    indices = convert_to_tensor(indices)
+    return jnp.unravel_index(indices, shape)
 
 
 @sparse.elementwise_unary(linear=True)
@@ -1212,7 +1212,7 @@ def vectorize(pyfunc, *, excluded=None, signature=None):
     return jnp.vectorize(pyfunc, excluded=excluded, signature=signature)
 
 
-def where(condition, x1, x2):
+def where(condition, x1=None, x2=None):
     return jnp.where(condition, x1, x2)
 
 
@@ -1357,5 +1357,5 @@ def argpartition(x, kth, axis=-1):
     return jnp.argpartition(x, kth, axis)
 
 
-def histogram(x, bins, range):
+def histogram(x, bins=10, range=None):
     return jnp.histogram(x, bins=bins, range=range)

keras/src/backend/numpy/core.py

@@ -343,14 +343,14 @@ def scatter_update(inputs, indices, updates):
     return inputs
 
 
-def slice(inputs, start_indices, lengths):
+def slice(inputs, start_indices, shape):
     # Validate inputs
-    assert len(start_indices) == len(lengths)
+    assert len(start_indices) == len(shape)
 
     # Generate list of indices arrays for each dimension
     indices = [
         np.arange(start, start + length)
-        for start, length in zip(start_indices, lengths)
+        for start, length in zip(start_indices, shape)
     ]
 
     # Use np.ix_ to create a multidimensional index array
@@ -407,8 +407,8 @@ def fori_loop(lower, upper, body_fun, init_val):
     return val
 
 
-def stop_gradient(x):
-    return x
+def stop_gradient(variable):
+    return variable
 
 
 def unstack(x, num=None, axis=0):

keras/src/backend/numpy/math.py

@@ -52,7 +52,7 @@ def segment_max(data, segment_ids, num_segments=None, sorted=False):
     )
 
 
-def top_k(x, k, sorted=False):
+def top_k(x, k, sorted=True):
     if sorted:
         # Take the k largest values.
         sorted_indices = np.argsort(x, axis=-1)[..., ::-1]

keras/src/backend/numpy/nn.py

@@ -125,11 +125,9 @@ def elu(x, alpha=1.0):
     )
 
 
-def selu(
-    x,
-    alpha=1.6732632423543772848170429916717,
-    scale=1.0507009873554804934193349852946,
-):
+def selu(x):
+    alpha = 1.6732632423543772848170429916717
+    scale = 1.0507009873554804934193349852946
     x = convert_to_tensor(x)
     return np.array(scale, x.dtype) * elu(x, alpha)
 
@@ -196,20 +194,20 @@ def threshold(x, threshold, default_value):
     return np.where(x > threshold, x, np.array(default_value, dtype=x.dtype))
 
 
-def softmax(x, axis=None):
+def softmax(x, axis=-1):
     exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
     return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
 
 
-def log_softmax(x, axis=None):
+def log_softmax(x, axis=-1):
     max_x = np.max(x, axis=axis, keepdims=True)
     logsumexp = np.log(np.exp(x - max_x).sum(axis=axis, keepdims=True))
     return x - max_x - logsumexp
 
 
-def sparsemax(logits, axis=-1):
+def sparsemax(x, axis=-1):
     # Sort logits along the specified axis in descending order
-    logits = convert_to_tensor(logits)
+    logits = convert_to_tensor(x)
     logits_sorted = -1.0 * np.sort(-1.0 * logits, axis=axis)
     logits_cumsum = np.cumsum(logits_sorted, axis=axis)
     r = np.arange(1, logits.shape[axis] + 1)
@@ -304,8 +302,8 @@ def max_pool(
 def average_pool(
     inputs,
     pool_size,
-    strides,
-    padding,
+    strides=None,
+    padding="valid",
     data_format=None,
 ):
     data_format = backend.standardize_data_format(data_format)
@@ -543,9 +541,11 @@ def conv_transpose(
     )
 
 
-def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
+def one_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with numpy backend")
+    if dtype is None:
+        dtype = "float32"
     x = convert_to_tensor(x)
     input_shape = x.shape
 
@@ -569,7 +569,7 @@ def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
     return categorical
 
 
-def multi_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
+def multi_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with numpy backend")
     x = convert_to_tensor(x)

keras/src/backend/numpy/numpy.py

@@ -173,7 +173,7 @@ def append(x1, x2, axis=None):
     return np.append(x1, x2, axis=axis)
 
 
-def arange(start, stop=None, step=None, dtype=None):
+def arange(start, stop=None, step=1, dtype=None):
     if dtype is None:
         dtypes_to_resolve = [
             getattr(start, "dtype", type(start)),
@@ -537,13 +537,13 @@ def digitize(x, bins):
     return np.digitize(x, bins).astype(np.int32)
 
 
-def dot(x, y):
-    x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
-    dtype = dtypes.result_type(x.dtype, y.dtype)
-    x = x.astype(dtype)
-    y = y.astype(dtype)
-    return np.dot(x, y)
+def dot(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    x1 = x1.astype(dtype)
+    x2 = x2.astype(dtype)
+    return np.dot(x1, x2)
 
 
 def empty(shape, dtype=None):
@@ -898,10 +898,10 @@ def ravel(x):
     return np.ravel(x)
 
 
-def unravel_index(x, shape):
-    dtype = dtypes.result_type(x.dtype)
+def unravel_index(indices, shape):
+    dtype = dtypes.result_type(indices.dtype)
     return tuple(
-        indices.astype(dtype) for indices in np.unravel_index(x, shape)
+        indices.astype(dtype) for indices in np.unravel_index(indices, shape)
     )
 
 
@@ -1114,7 +1114,7 @@ def vectorize(pyfunc, *, excluded=None, signature=None):
     return np.vectorize(pyfunc, excluded=excluded, signature=signature)
 
 
-def where(condition, x1, x2):
+def where(condition, x1=None, x2=None):
     if x1 is not None and x2 is not None:
         if not isinstance(x1, (int, float)):
             x1 = convert_to_tensor(x1)
@@ -1283,5 +1283,5 @@ def argpartition(x, kth, axis=-1):
     return np.argpartition(x, kth, axis).astype("int32")
 
 
-def histogram(x, bins, range):
+def histogram(x, bins=10, range=None):
     return np.histogram(x, bins=bins, range=range)

keras/src/backend/openvino/core.py

@@ -565,21 +565,21 @@ def scatter_update(inputs, indices, updates):
     )
 
 
-def slice(inputs, start_indices, lengths):
+def slice(inputs, start_indices, shape):
     inputs = get_ov_output(inputs)
     assert isinstance(start_indices, tuple), (
         "`slice` is not supported by openvino backend"
-        " for `start_indices` of type {}".format(type(lengths))
+        " for `start_indices` of type {}".format(type(shape))
     )
-    assert isinstance(lengths, tuple), (
+    assert isinstance(shape, tuple), (
         "`slice` is not supported by openvino backend"
-        " for `lengths` of type {}".format(type(lengths))
+        " for `lengths` of type {}".format(type(shape))
     )
 
     axes = []
     start = []
     stop = []
-    for idx, length in enumerate(lengths):
+    for idx, length in enumerate(shape):
         if length is not None and length >= 0:
             axes.append(idx)
             start.append(start_indices[idx])
@@ -621,8 +621,8 @@ def fori_loop(lower, upper, body_fun, init_val):
     )
 
 
-def stop_gradient(x):
-    return x
+def stop_gradient(variable):
+    return variable
 
 
 def unstack(x, num=None, axis=0):

keras/src/backend/openvino/image.py

@@ -1,4 +1,4 @@
-def rgb_to_grayscale(image, data_format="channels_last"):
+def rgb_to_grayscale(images, data_format=None):
     raise NotImplementedError(
         "`rgb_to_grayscale` is not supported with openvino backend"
     )
@@ -19,20 +19,20 @@ def resize(
 
 
 def affine_transform(
-    image,
+    images,
     transform,
     interpolation="bilinear",
     fill_mode="constant",
     fill_value=0,
-    data_format="channels_last",
+    data_format=None,
 ):
     raise NotImplementedError(
         "`affine_transform` is not supported with openvino backend"
     )
 
 
 def map_coordinates(
-    input, coordinates, order, fill_mode="constant", fill_value=0.0
+    inputs, coordinates, order, fill_mode="constant", fill_value=0
 ):
     raise NotImplementedError(
         "`map_coordinates` is not supported with openvino backend"