Repository: keras-team/keras (python)

Issue:
Numpy Ops function nonzero(x) appers to be missing check for symbolic tensors
In updating  code from Keras 2 to 3, we noticed that nonzero function continues to throw errors for use of KerasTensor in TF functions, even when run though tf.keras.ops

Digging into the source, it appears that this function does not receive the check for any_symbolic_tensors(), and thus no instantiation of the NonZero() class. In turn failing when used with a KerasTensor

https://github.com/keras-team/keras/blob/42a1535ed7d3d75711a11d295f58a2dc9a59fdae/keras/ops/numpy.py#L3976

Pull request description:
Add symbolic shape to nonzero
Fix #19407.

Patch:
diff --git a/keras/backend/common/variables.py b/keras/backend/common/variables.py
index 2549b0a79..e69a01e2e 100644
--- a/keras/backend/common/variables.py
+++ b/keras/backend/common/variables.py
@@ -12,6 +12,71 @@ from keras.utils.naming import auto_name
 
 
 class KerasVariable:
+    """Represents a backend-agnostic variable in Keras.
+
+    A `Variable` acts as a container for state. It holds a tensor value and can
+    be updated. With the JAX backend, variables are used to implement
+    "functionalization", the pattern of lifting stateful operations out of
+    a piece of computation to turn it into a stateless function.
+
+    Args:
+        initializer: Initial value or callable for initialization.
+            If a callable is used, it should take the arguments
+            `shape` and `dtype`.
+        shape: Optional. Tuple for the variable's shape. Required if `initializer`
+            is a callable.
+        dtype: Optional. Data type of the variable. Defaults to the global float
+            dtype type (`"float32"` if never configured).
+        trainable: Optional. Boolean indicating if variable is trainable.
+            Defaults to `True`.
+        name: Optional. A unique name for the variable. Automatically generated
+            if not set.
+
+    Attributes:
+        name: The name of the variable (string).
+        path: The path of the variable within the Keras model or layer (string).
+        dtype: The data type of the variable (string).
+        shape: The shape of the variable (tuple of integers).
+        ndim: The number of dimensions of the variable (integer).
+        trainable: Whether the variable is trainable (boolean).
+        value: The current value of the variable (NumPy array or tensor).
+
+    Examples:
+
+    **Initializing a `Variable` with a NumPy array:**
+
+    ```python
+    import numpy as np
+    import keras
+    initial_array = np.ones((3, 3))
+    variable_from_array = keras.Variable(initializer=initial_array)
+    ```
+
+    **Using a Keras initializer to create a `Variable`:**
+
+    ```python
+    from keras.initializers import Ones
+    variable_from_initializer = keras.Variable(
+        initializer=Ones(), shape=(3, 3), dtype="float32"
+    )
+    ```
+
+    **Updating the value of a `Variable`:**
+
+    ```python
+    new_value = np.zeros((3, 3), dtype="float32")
+    variable_from_array.assign(new_value)
+    ```
+
+    **Marking a `Variable` as non-trainable:**
+
+    ```python
+    non_trainable_variable = keras.Variable(
+        initializer=np.ones((3, 3), dtype="float32"), trainable=False
+    )
+    ```
+    """
+
     def __init__(
         self,
         initializer,
diff --git a/keras/ops/numpy.py b/keras/ops/numpy.py
index 25f1809d6..dffe381f2 100644
--- a/keras/ops/numpy.py
+++ b/keras/ops/numpy.py
@@ -3992,6 +3992,9 @@ class Nonzero(Operation):
     def call(self, x):
         return backend.numpy.nonzero(x)
 
+    def compute_output_spec(self, x):
+        return KerasTensor([None] * len(x.shape))
+
 
 @keras_export(["keras.ops.nonzero", "keras.ops.numpy.nonzero"])
 def nonzero(x):
@@ -4003,6 +4006,8 @@ def nonzero(x):
     Returns:
         Indices of elements that are non-zero.
     """
+    if any_symbolic_tensors((x,)):
+        return Nonzero().symbolic_call(x)
     return backend.numpy.nonzero(x)
 
 


Test patch:
diff --git a/keras/ops/numpy_test.py b/keras/ops/numpy_test.py
index 42a451c2c..07f05aa20 100644
--- a/keras/ops/numpy_test.py
+++ b/keras/ops/numpy_test.py
@@ -1311,6 +1311,10 @@ class NumpyOneInputOpsDynamicShapeTest(testing.TestCase):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.ndim(x).shape, (2,))
 
+    def test_nonzero(self):
+        x = KerasTensor((None, 5, 6))
+        self.assertEqual(knp.nonzero(x).shape, (None, None, None))
+
     def test_ones_like(self):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.ones_like(x).shape, (None, 3))