Fix operations which had static parameter passed as dynamic parameters. (#21374)

Switched many parameters in ops from being dynamic (passed to `call` method) to being static (passed to `__init__`). These parameters should be static because they cannot be tensors. The most commonly changed parameters were `shape` and `dtype`.

Standardized the way `dtype` parameters are handled, in particular when they can be `None`.

Removed unused op classes: `Empty`, `Eye`, `Identity`, `Ones`, `Tri`, `Zeros`. These ops only have static parameters and no tensor inputs, they therefore cannot be part of a functional graph. They were missing the standard `if any_symbolic_tensors(...)`, which is why the class was dead code and only instantiated explicitly by unit tests.

Fixed `arange` and `full` ops. The code was missing the standard `if any_symbolic_tensors(...)`, which means that adding them to a functional graph would fail before and the op class was dead code until now.

Author: hertschuh

keras/src/ops/core.py

@@ -78,24 +78,30 @@ def map(f, xs):
 
 
 class Scan(Operation):
-    def __init__(self, reverse=False, unroll=1):
+    def __init__(self, length=None, reverse=False, unroll=1):
         super().__init__()
+        self.length = length
         self.reverse = reverse
         self.unroll = unroll
 
-    def call(self, f, init, xs=None, length=None):
+    def call(self, f, init, xs=None):
         return backend.core.scan(
-            f, init, xs, length, reverse=self.reverse, unroll=self.unroll
+            f,
+            init,
+            xs,
+            length=self.length,
+            reverse=self.reverse,
+            unroll=self.unroll,
         )
 
-    def compute_output_spec(self, f, init, xs=None, length=None):
+    def compute_output_spec(self, f, init, xs=None):
         if xs is None:
-            n = int(length)
+            n = int(self.length)
             x = None
         else:
             n = (
-                int(length)
-                if length is not None
+                int(self.length)
+                if self.length is not None
                 else tree.flatten(xs)[0].shape[0]
             )
             x = xs[0]
@@ -176,9 +182,9 @@ def scan(f, init, xs, length=None):
     [1, 3, 6, 10, 15]
     """
     if any_symbolic_tensors((init, xs)):
-        return Scan(reverse=reverse, unroll=unroll).symbolic_call(
-            f, init, xs, length
-        )
+        return Scan(
+            length=length, reverse=reverse, unroll=unroll
+        ).symbolic_call(f, init, xs)
     return backend.core.scan(
         f, init, xs, length, reverse=reverse, unroll=unroll
     )
@@ -283,11 +289,15 @@ def associative_scan(f, elems, reverse=False, axis=0):
 
 
 class Scatter(Operation):
-    def call(self, indices, values, shape):
-        return backend.core.scatter(indices, values, shape)
+    def __init__(self, shape):
+        super().__init__()
+        self.shape = shape
 
-    def compute_output_spec(self, indices, values, shape):
-        return KerasTensor(shape, dtype=values.dtype)
+    def call(self, indices, values):
+        return backend.core.scatter(indices, values, self.shape)
+
+    def compute_output_spec(self, indices, values):
+        return KerasTensor(self.shape, dtype=values.dtype)
 
 
 @keras_export("keras.ops.scatter")
@@ -316,8 +326,8 @@ def scatter(indices, values, shape):
     array([[0., 1.],
            [0., 1.]])
     """
-    if any_symbolic_tensors((indices, values, shape)):
-        return Scatter().symbolic_call(indices, values, shape)
+    if any_symbolic_tensors((indices, values)):
+        return Scatter(shape=shape).symbolic_call(indices, values)
     return backend.core.scatter(indices, values, shape)
 
 
@@ -382,11 +392,15 @@ def scatter_update(inputs, indices, updates):
 
 
 class Slice(Operation):
-    def call(self, inputs, start_indices, shape):
-        return backend.core.slice(inputs, start_indices, shape)
+    def __init__(self, shape):
+        super().__init__()
+        self.shape = shape
 
-    def compute_output_spec(self, inputs, start_indices, shape):
-        return KerasTensor(shape, dtype=inputs.dtype)
+    def call(self, inputs, start_indices):
+        return backend.core.slice(inputs, start_indices, self.shape)
+
+    def compute_output_spec(self, inputs, start_indices):
+        return KerasTensor(self.shape, dtype=inputs.dtype)
 
 
 @keras_export("keras.ops.slice")
@@ -415,8 +429,8 @@ def slice(inputs, start_indices, shape):
     Returns:
         A tensor, has the same shape and dtype as `inputs`.
     """
-    if any_symbolic_tensors((inputs, start_indices, shape)):
-        return Slice().symbolic_call(inputs, start_indices, shape)
+    if any_symbolic_tensors((inputs, start_indices)):
+        return Slice(shape=shape).symbolic_call(inputs, start_indices)
     return backend.core.slice(inputs, start_indices, shape)
 
 
@@ -916,7 +930,7 @@ def get_dtype_min_max(dtype):
 class ConvertToTensor(Operation):
     def __init__(self, dtype=None, sparse=None, ragged=None):
         super().__init__()
-        self.dtype = backend.standardize_dtype(dtype)
+        self.dtype = None if dtype is None else backend.standardize_dtype(dtype)
         self.sparse = sparse
         self.ragged = ragged
 
@@ -926,7 +940,11 @@ def call(self, x):
         )
 
     def compute_output_spec(self, x):
-        dtype = x.dtype if self.dtype is None else self.dtype
+        dtype = (
+            backend.standardize_dtype(x.dtype)
+            if self.dtype is None
+            else self.dtype
+        )
         sparse = (
             False if self.sparse is not None and not self.sparse else x.sparse
         )

keras/src/ops/core_test.py

@@ -1010,7 +1010,7 @@ def cumsum(carry, xs):
         init = np.array(0, dtype="float32")
         xs = np.array([1, 2, 3, 4, 10, 20], dtype="float32")
         scan_op = core.Scan()
-        carry, result = scan_op.call(cumsum, init, xs, None)
+        carry, result = scan_op.call(cumsum, init, xs)
         self.assertAllClose(carry, 40.0)
         self.assertAllClose(result, ops.cumsum(xs))
 
@@ -1025,8 +1025,8 @@ def test_scatter_basic_call(self):
         indices = np.array([[1, 0], [0, 1]])
         values = np.array([10, 20])
         shape = (2, 2)
-        scatter = core.Scatter()
-        result = scatter.call(indices, values, shape)
+        scatter = core.Scatter(shape)
+        result = scatter.call(indices, values)
         expected_output = np.array([[0, 20], [10, 0]])
         self.assertAllClose(core.convert_to_numpy(result), expected_output)
 
@@ -1043,17 +1043,17 @@ def test_slice_basic_call(self):
         inputs = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
         start_indices = np.array([1, 1])
         shape = (2, 2)
-        slice_op = core.Slice()
-        result = slice_op.call(inputs, start_indices, shape)
+        slice_op = core.Slice(shape)
+        result = slice_op.call(inputs, start_indices)
         expected_output = np.array([[5, 6], [8, 9]])
         self.assertAllClose(core.convert_to_numpy(result), expected_output)
 
     def test_slice_compute_output_spec(self):
         inputs = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
         start_indices = np.array([1, 1])
         shape = (2, 2)
-        slice_op = core.Slice()
-        output_spec = slice_op.compute_output_spec(inputs, start_indices, shape)
+        slice_op = core.Slice(shape)
+        output_spec = slice_op.compute_output_spec(inputs, start_indices)
         self.assertEqual(output_spec.shape, shape)
         self.assertEqual(output_spec.dtype, inputs.dtype)
 

keras/src/ops/nn.py

@@ -1693,7 +1693,7 @@ def __init__(self, num_classes, axis=-1, dtype=None, sparse=False):
         super().__init__()
         self.num_classes = num_classes
         self.axis = axis
-        self.dtype = dtype or backend.floatx()
+        self.dtype = backend.standardize_dtype(dtype)
         self.sparse = sparse
 
     def call(self, x):
@@ -2580,9 +2580,13 @@ def psnr(
 
 
 class DotProductAttention(Operation):
-    def __init__(self, is_causal=False):
+    def __init__(
+        self, is_causal=False, flash_attention=None, attn_logits_soft_cap=None
+    ):
         super().__init__()
         self.is_causal = is_causal
+        self.flash_attention = flash_attention
+        self.attn_logits_soft_cap = attn_logits_soft_cap
 
     def call(
         self,
@@ -2592,8 +2596,6 @@ def call(
         bias=None,
         mask=None,
         scale=None,
-        flash_attention=None,
-        attn_logits_soft_cap=None,
     ):
         return backend.nn.dot_product_attention(
             query,
@@ -2603,8 +2605,8 @@ def call(
             mask=mask,
             scale=scale,
             is_causal=self.is_causal,
-            flash_attention=flash_attention,
-            attn_logits_soft_cap=attn_logits_soft_cap,
+            flash_attention=self.flash_attention,
+            attn_logits_soft_cap=self.attn_logits_soft_cap,
         )
 
     def compute_output_spec(
@@ -2615,8 +2617,6 @@ def compute_output_spec(
         bias=None,
         mask=None,
         scale=None,
-        flash_attention=None,
-        attn_logits_soft_cap=None,
     ):
         return KerasTensor(query.shape, dtype=query.dtype)
 
@@ -2703,15 +2703,17 @@ def dot_product_attention(
             )
 
     if any_symbolic_tensors((query, key, value)):
-        return DotProductAttention(is_causal=is_causal).symbolic_call(
+        return DotProductAttention(
+            is_causal=is_causal,
+            flash_attention=flash_attention,
+            attn_logits_soft_cap=attn_logits_soft_cap,
+        ).symbolic_call(
             query,
             key,
             value,
             bias=bias,
             mask=mask,
             scale=scale,
-            flash_attention=flash_attention,
-            attn_logits_soft_cap=attn_logits_soft_cap,
         )
     return backend.nn.dot_product_attention(
         query,