Repository: keras-team/keras (Python)
Task: Bug Fix (2024-03-13 07:57:31+00:00)

Issue:
`keras.ops.softmax` errors out when used in a TensorFlow compiled function
## MRE

```python
import keras
from keras import ops

class SoftmaxLayer(keras.Layer):
  def call(self, x):
    return ops.softmax(x, axis=-1)

class Model(keras.Model):
  def __init__(self):
    x = keras.Input(shape=(None,))
    y = SoftmaxLayer()(x)
    super().__init__(inputs=x, outputs=y)

model = Model()  # Error
```

## Additional Details

The regression was introduced in [d5a4521](https://github.com/keras-team/keras/commit/d5a452155a415d5fcbe568eb2a8441f64e57aa90)

Discovered by KerasCV test run with `keras-nightly`: https://github.com/keras-team/keras-cv/actions/runs/8259807616/job/22594330565


Discussion:

Patch:
diff --git a/keras/ops/nn.py b/keras/ops/nn.py
--- a/keras/ops/nn.py
+++ b/keras/ops/nn.py
@@ -538,10 +538,13 @@ def softmax(x, axis=-1):
     array([0.09003057, 0.24472847, 0.66524096], shape=(3,), dtype=float64)
 
     """
-    if isinstance(axis, int) and backend.shape(x)[axis] == 1:
+    # Don't use `backend.shape` since TensorFlow returns
+    # symbolic tensors for unknown shape which can trigger
+    # an error in TensorFlow graph execution.
+    if isinstance(axis, int) and x.shape[axis] == 1:
         warnings.warn(
             f"You are using a softmax over axis {axis} "
-            f"of a tensor of shape {backend.shape(x)}. This axis "
+            f"of a tensor of shape {x.shape}. This axis "
             "has size 1. The softmax operation will always return "
             "the value 1, which is likely not what you intended. "
             "Did you mean to use a sigmoid instead?"


Test patch:
diff --git a/keras/ops/nn_test.py b/keras/ops/nn_test.py
--- a/keras/ops/nn_test.py
+++ b/keras/ops/nn_test.py
@@ -2,10 +2,12 @@
 import pytest
 from absl.testing import parameterized
 
+import keras
 from keras import backend
 from keras import layers
 from keras import losses
 from keras import models
+from keras import ops
 from keras import testing
 from keras.backend.common import standardize_dtype
 from keras.backend.common.keras_tensor import KerasTensor
@@ -84,6 +86,22 @@ def test_softmax(self):
         self.assertEqual(knn.softmax(x, axis=1).shape, (None, 2, 3))
         self.assertEqual(knn.softmax(x, axis=-1).shape, (None, 2, 3))
 
+    def test_softmax_in_graph(self):
+        class SoftmaxLayer(keras.Layer):
+            def call(self, x):
+                return ops.softmax(x, axis=-1)
+
+        class Model(keras.Model):
+            def __init__(self):
+                x = keras.Input(shape=(None,))
+                y = SoftmaxLayer()(x)
+                super().__init__(inputs=x, outputs=y)
+
+        # Make sure Keras is able to compile the model graph
+        model = Model()
+        x = ops.array([[1.0, 2.0, 3.0, 4.0]])
+        model.predict(x)
+
     def test_log_softmax(self):
         x = KerasTensor([None, 2, 3])
         self.assertEqual(knn.log_softmax(x).shape, (None, 2, 3))