Repository: keras-team/keras

Issue:
keras.ops.tril doesn't support dynamic shape with tensorflow backend
keras.ops.tril doesn't support tensor with unknown dimension with tensorflow backend.
it is OK with torch backend.
could you update keras.ops.tril implement with tensorflow ops to support input with unknown dimension?

```python
#!/usr/bin/python3

from os import environ
environ['KERAS_BACKEND'] = 'tensorflow'  # FAIL
#environ['KERAS_BACKEND'] = 'torch' # OK
import keras as K

def Test():
  inputs = K.Input((None, 100)) # inputs.shape = (batch x seq_len x hidden)
  attn = K.layers.Lambda(lambda x: K.ops.tril(K.ops.ones((K.ops.shape(x)[1],K.ops.shape(x)[1]))))(inputs)
  return K.Model(inputs = inputs, outputs = attn)

test = Test()
import numpy as np
print(test(np.random.normal(size = (1,10,100))).shape)
```

I find the following tensorflow implement support dimension with unknown dimension

```python
#!/usr/bin/python3

import tensorflow as tf

def Test():
  inputs = tf.keras.Input((None, 100))
  outputs = tf.keras.layers.Lambda(lambda x: tf.linalg.band_part(tf.ones((tf.shape(x)[1], tf.shape(x)[1])), -1, 0))(inputs)
  return tf.keras.Model(inputs = inputs, outputs = outputs)

test = Test()
import numpy as np
print(test(np.random.normal(size = (1,10,100))))
```


Discussion:
Hi @breadbread1984 ,

I have replicated the issue with TF backend.Attached [gist-tf](https://colab.sandbox.google.com/gist/SuryanarayanaY/73047a12f402fb7b50a9e4335c99e65e/18890_tf-backend.ipynb) for reference.

It works with Torch as reported.[gist-torch](https://colab.sandbox.google.com/gist/SuryanarayanaY/26e1503ce61fba33edc1486e1bf502ae/18890_torch-backend.ipynb)

Resolving patch:
diff --git a/keras/backend/tensorflow/numpy.py b/keras/backend/tensorflow/numpy.py
--- a/keras/backend/tensorflow/numpy.py
+++ b/keras/backend/tensorflow/numpy.py
@@ -1462,20 +1462,26 @@ def tri(N, M=None, k=0, dtype=None):
 
 def tril(x, k=0):
     x = convert_to_tensor(x)
-    # TODO: tfnp.tril doesn't support bool
-    if standardize_dtype(x.dtype) == "bool":
-        x = tf.cast(x, "uint8")
-        return tf.cast(tfnp.tril(x, k=k), "bool")
-    return tfnp.tril(x, k=k)
+    if k >= 0:
+        return tf.linalg.band_part(x, -1, k)
+
+    # deal with negative k using mask
+    k = -k - 1
+    mask = tf.ones_like(x, dtype="bool")
+    mask = tf.logical_not(tf.linalg.band_part(mask, k, -1))
+    return tf.where(mask, x, tf.constant(0, x.dtype))
 
 
 def triu(x, k=0):
     x = convert_to_tensor(x)
-    # TODO: tfnp.triu doesn't support bool
-    if standardize_dtype(x.dtype) == "bool":
-        x = tf.cast(x, "uint8")
-        return tf.cast(tfnp.tril(x, k=k), "bool")
-    return tfnp.triu(x, k=k)
+    if k >= 0:
+        return tf.linalg.band_part(x, k, -1)
+
+    # deal with negative k using mask
+    k = -k
+    mask = tf.ones_like(x, dtype="bool")
+    mask = tf.logical_not(tf.linalg.band_part(mask, k, -1))
+    return tf.where(mask, tf.constant(0, x.dtype), x)
 
 
 def vdot(x1, x2):


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
@@ -6,6 +6,7 @@
 from absl.testing import parameterized
 from tensorflow.python.ops.numpy_ops import np_config
 
+import keras
 from keras import backend
 from keras import testing
 from keras.backend.common import standardize_dtype
@@ -3840,12 +3841,54 @@ def test_tril(self):
         self.assertAllClose(knp.tril(x, -1), np.tril(x, -1))
         self.assertAllClose(knp.Tril(-1)(x), np.tril(x, -1))
 
+    def test_tril_in_layer(self):
+        # https://github.com/keras-team/keras/issues/18890
+        x = keras.Input((None, 3))
+        y1 = keras.layers.Lambda(
+            lambda x: keras.ops.tril(
+                keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1]))
+            )
+        )(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,
+            )
+        )(x)
+        model = keras.Model(x, [y1, y2])
+
+        result = model(np.ones((1, 2, 3), "float32"))
+        self.assertAllClose(
+            result, [np.tril(np.ones((2, 2))), np.tril(np.ones((2, 2)), k=-1)]
+        )
+
     def test_triu(self):
         x = np.arange(24).reshape([1, 2, 3, 4])
         self.assertAllClose(knp.triu(x), np.triu(x))
         self.assertAllClose(knp.triu(x, -1), np.triu(x, -1))
         self.assertAllClose(knp.Triu(-1)(x), np.triu(x, -1))
 
+    def test_triu_in_layer(self):
+        # https://github.com/keras-team/keras/issues/18890
+        x = keras.Input((None, 3))
+        y1 = keras.layers.Lambda(
+            lambda x: keras.ops.triu(
+                keras.ops.ones((keras.ops.shape(x)[1], keras.ops.shape(x)[1]))
+            )
+        )(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,
+            )
+        )(x)
+        model = keras.Model(x, [y1, y2])
+
+        result = model(np.ones((1, 2, 3), "float32"))
+        self.assertAllClose(
+            result, [np.triu(np.ones((2, 2))), np.triu(np.ones((2, 2)), k=-1)]
+        )
+
     def test_vstack(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [6, 5, 4]])