Repository: keras-team/keras (Python)
Task type: Bug Fix

Issue:
Ops inconsistency with tensorflow for tril and triu
`ops.tril` and `ops.triu` allow specifying a negative diagonal. For compiled runs, we use a python conditional instead of cond to check the sign of the diagonal which breaks. This is tensorflow specific, other backends allow negative diagonals.

```
../miniconda3/envs/keras-nlp-cpu/lib/python3.10/site-packages/keras/src/backend/tensorflow/numpy.py:2145: in triu
    if k <= 0:
../miniconda3/envs/keras-nlp-cpu/lib/python3.10/site-packages/tensorflow/python/framework/tensor.py:660: in __bool__
    self._disallow_bool_casting()
../miniconda3/envs/keras-nlp-cpu/lib/python3.10/site-packages/tensorflow/python/framework/tensor.py:316: in _disallow_bool_casting
    self._disallow("Using a symbolic `tf.Tensor` as a Python `bool`")

E     tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: Using a symbolic `tf.Tensor` as a Python `bool` is not allowed. You can attempt the following resolutions to the problem: If you are running in Graph mode, use Eager execution mode or decorate this function with @tf.function. If you are using AutoGraph, you can try decorating this function with @tf.function. If that does not work, then you may be using an unsupported feature or your source code may not be visible to AutoGraph. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code for more information.

../miniconda3/envs/keras-nlp-cpu/lib/python3.10/site-packages/tensorflow/python/framework/tensor.py:303: OperatorNotAllowedInGraphError
```


Discussion:

Solution patch:
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
--- a/keras/src/backend/tensorflow/numpy.py
+++ b/keras/src/backend/tensorflow/numpy.py
@@ -2126,33 +2126,31 @@ def tri(N, M=None, k=0, dtype=None):
 def tril(x, k=0):
     x = convert_to_tensor(x)
 
-    if k >= 0:
-        return tf.linalg.band_part(x, -1, k)
+    def _negative_k_branch():
+        shape = tf.shape(x)
+        rows, cols = shape[-2], shape[-1]
+        i, j = tf.meshgrid(tf.range(rows), tf.range(cols), indexing="ij")
+        mask = i >= j - k
+        return tf.where(tf.broadcast_to(mask, shape), x, tf.zeros_like(x))
 
-    shape = tf.shape(x)
-    rows, cols = shape[-2], shape[-1]
-
-    i, j = tf.meshgrid(tf.range(rows), tf.range(cols), indexing="ij")
-
-    mask = i >= j - k
-
-    return tf.where(tf.broadcast_to(mask, shape), x, tf.zeros_like(x))
+    return tf.cond(
+        k >= 0, lambda: tf.linalg.band_part(x, -1, k), _negative_k_branch
+    )
 
 
 def triu(x, k=0):
     x = convert_to_tensor(x)
 
-    if k <= 0:
-        return tf.linalg.band_part(x, -k, -1)
+    def _positive_k_branch():
+        shape = tf.shape(x)
+        rows, cols = shape[-2], shape[-1]
+        i, j = tf.meshgrid(tf.range(rows), tf.range(cols), indexing="ij")
+        mask = i <= j - k
+        return tf.where(tf.broadcast_to(mask, shape), x, tf.zeros_like(x))
 
-    shape = tf.shape(x)
-    rows, cols = shape[-2], shape[-1]
-
-    i, j = tf.meshgrid(tf.range(rows), tf.range(cols), indexing="ij")
-
-    mask = i <= j - k
-
-    return tf.where(tf.broadcast_to(mask, shape), x, tf.zeros_like(x))
+    return tf.cond(
+        k <= 0, lambda: tf.linalg.band_part(x, -k, -1), _positive_k_branch
+    )
 
 
 def vdot(x1, x2):


Test patch:
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -4168,13 +4168,15 @@ def test_tril_in_layer(self):
         y1 = keras.layers.Lambda(
             lambda x: keras.ops.tril(
                 keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1]))
-            )
+            ),
+            output_shape=(None, None, 3),
         )(x)
         y2 = keras.layers.Lambda(
             lambda x: keras.ops.tril(
                 keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1])),
                 k=-1,
-            )
+            ),
+            output_shape=(None, None, 3),
         )(x)
         model = keras.Model(x, [y1, y2])
 
@@ -4183,6 +4185,24 @@ def test_tril_in_layer(self):
             result, [np.tril(np.ones((2, 2))), np.tril(np.ones((2, 2)), k=-1)]
         )
 
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="Only test tensorflow backend",
+    )
+    def test_tril_with_jit_in_tf(self):
+        import tensorflow as tf
+
+        x = knp.reshape(knp.arange(24), [1, 2, 3, 4])
+        k = knp.array(0)
+        x_np = np.reshape(np.arange(24), [1, 2, 3, 4])
+        k_np = np.array(0)
+
+        @tf.function(jit_compile=True)
+        def fn(x, k):
+            return knp.tril(x, k=k)
+
+        self.assertAllClose(fn(x, k), np.tril(x_np, k_np))
+
     def test_triu(self):
         x = np.arange(24).reshape([1, 2, 3, 4])
         self.assertAllClose(knp.triu(x), np.triu(x))
@@ -4200,13 +4220,15 @@ def test_triu_in_layer(self):
         y1 = keras.layers.Lambda(
             lambda x: keras.ops.triu(
                 keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1]))
-            )
+            ),
+            output_shape=(None, None, 3),
         )(x)
         y2 = keras.layers.Lambda(
             lambda x: keras.ops.triu(
                 keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1])),
                 k=-1,
-            )
+            ),
+            output_shape=(None, None, 3),
         )(x)
         model = keras.Model(x, [y1, y2])
 
@@ -4215,6 +4237,24 @@ def test_triu_in_layer(self):
             result, [np.triu(np.ones((2, 2))), np.triu(np.ones((2, 2)), k=-1)]
         )
 
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="Only test tensorflow backend",
+    )
+    def test_triu_with_jit_in_tf(self):
+        import tensorflow as tf
+
+        x = knp.reshape(knp.arange(24), [1, 2, 3, 4])
+        k = knp.array(0)
+        x_np = np.reshape(np.arange(24), [1, 2, 3, 4])
+        k_np = np.array(0)
+
+        @tf.function(jit_compile=True)
+        def fn(x, k):
+            return knp.triu(x, k=k)
+
+        self.assertAllClose(fn(x, k), np.triu(x_np, k_np))
+
     def test_vstack(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [6, 5, 4]])