Repository: keras-team/keras (python)

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
```

Pull request description:
Fix `ops.tril` and `ops.triu` for tensorflow backend
Fix #19919

The tests for `tf.function(jit_compile=True)` have been included.

Patch:
diff --git a/keras/src/backend/common/keras_tensor.py b/keras/src/backend/common/keras_tensor.py
index 2876a57cf..d0ff42516 100644
--- a/keras/src/backend/common/keras_tensor.py
+++ b/keras/src/backend/common/keras_tensor.py
@@ -37,12 +37,45 @@ class KerasTensor:
     ):
         from keras.src import backend
 
-        self.shape = backend.standardize_shape(shape)
-        self.dtype = backend.standardize_dtype(dtype)
-        self.sparse = sparse
+        self._shape = backend.standardize_shape(shape)
+        self._dtype = backend.standardize_dtype(dtype)
+        self._sparse = bool(sparse)
         self.name = name or auto_name(self.__class__.__name__)
         self.record_history = record_history
 
+    @property
+    def shape(self):
+        return self._shape
+
+    @shape.setter
+    def shape(self, value):
+        raise AttributeError(
+            f"The shape of {self.__class__.__name__} is immutable. One should "
+            "create a new instance of KerasTensor for this."
+        )
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @dtype.setter
+    def dtype(self, value):
+        raise AttributeError(
+            f"The dtype of {self.__class__.__name__} is immutable. One should "
+            "create a new instance of KerasTensor for this."
+        )
+
+    @property
+    def sparse(self):
+        return self._sparse
+
+    @sparse.setter
+    def sparse(self, value):
+        raise AttributeError(
+            f"The sparse of {self.__class__.__name__} is immutable. One should "
+            "create a new instance of KerasTensor for this."
+        )
+
     @property
     def ndim(self):
         return len(self.shape)
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
index 25559efa1..8a7db1650 100644
--- 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):
diff --git a/keras/src/dtype_policies/dtype_policy.py b/keras/src/dtype_policies/dtype_policy.py
index cf7231d38..9d37ac49f 100644
--- a/keras/src/dtype_policies/dtype_policy.py
+++ b/keras/src/dtype_policies/dtype_policy.py
@@ -1,4 +1,5 @@
 from keras.src import backend
+from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.backend.common import global_state
 
@@ -161,7 +162,7 @@ class DTypePolicy:
             return x
         elif backend.is_keras_tensor(x):
             if self._should_cast(x, autocast, dtype):
-                x.dtype = dtype
+                x = ops.cast(x, dtype=dtype)
             return x
         elif hasattr(x, "__array__"):
             try:
diff --git a/keras/src/ops/core.py b/keras/src/ops/core.py
index 43a65ab92..b54947a36 100644
--- a/keras/src/ops/core.py
+++ b/keras/src/ops/core.py
@@ -39,8 +39,9 @@ class Map(Operation):
         y = backend.compute_output_spec(f, x)
 
         def append_batch_axis(x):
-            x.shape = (n,) + x.shape
-            return x
+            return KerasTensor(
+                shape=(n,) + x.shape, dtype=x.dtype, sparse=x.sparse
+            )
 
         y = tree.map_structure(append_batch_axis, y)
         return y
@@ -114,9 +115,9 @@ class Scan(Operation):
             )
             x = xs[0]
 
-        carry_spec, y_spec = backend.compute_output_spec(f, init, x)
-        y_spec.shape = (n,) + y_spec.shape
-        return carry_spec, y_spec
+        carry, y = backend.compute_output_spec(f, init, x)
+        y = KerasTensor(shape=(n,) + y.shape, dtype=y.dtype, sparse=y.sparse)
+        return carry, y
 
 
 @keras_export("keras.ops.scan")
diff --git a/keras/src/ops/function.py b/keras/src/ops/function.py
index f34ea3617..3e5daf035 100644
--- a/keras/src/ops/function.py
+++ b/keras/src/ops/function.py
@@ -127,8 +127,8 @@ class Function(Operation):
         # Ensure that dtype and sparse settings are the same as self._inputs,
         # because we only care about the shape in this function.
         for x, x_ref in zip(tree.flatten(input_shape_struct), self._inputs):
-            x.dtype = x_ref.dtype
-            x.sparse = x_ref.sparse
+            x._dtype = x_ref.dtype
+            x._sparse = x_ref.sparse
         output_spec = self.compute_output_spec(input_shape_struct)
         return tree.map_structure(lambda x: x.shape, output_spec)
 


Test patch:
diff --git a/keras/src/backend/common/keras_tensor_test.py b/keras/src/backend/common/keras_tensor_test.py
index ca6391024..fee822233 100644
--- a/keras/src/backend/common/keras_tensor_test.py
+++ b/keras/src/backend/common/keras_tensor_test.py
@@ -17,6 +17,20 @@ class KerasTensorTest(testing.TestCase):
         self.assertEqual(x.shape, (3,))
         self.assertEqual(x.sparse, True)
 
+        # Raise error if trying to set attributes
+        with self.assertRaisesRegex(
+            AttributeError, "The shape of KerasTensor is immutable."
+        ):
+            x.shape = [3, 2]
+        with self.assertRaisesRegex(
+            AttributeError, "The dtype of KerasTensor is immutable."
+        ):
+            x.dtype = "int32"
+        with self.assertRaisesRegex(
+            AttributeError, "The sparse of KerasTensor is immutable."
+        ):
+            x.sparse = False
+
     def test_numpy_methods(self):
         x = keras_tensor.KerasTensor(shape=(3, 2), dtype="float32")
 
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
index 2bfe4a637..3a142a153 100644
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -4168,13 +4168,15 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
         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 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
             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 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
         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 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
             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]])