Repository: keras-team/keras (Python)
Task: Bug Fix (2024-04-09 17:23:58+00:00)

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


Discussion:

Patch:
diff --git a/keras/ops/numpy.py b/keras/ops/numpy.py
--- 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
--- a/keras/ops/numpy_test.py
+++ b/keras/ops/numpy_test.py
@@ -1311,6 +1311,10 @@ def test_ndim(self):
         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))