Fix the shape of KerasTensor for glu. (#21696)

Author: james77777778

keras/src/backend/numpy/nn.py

@@ -164,13 +164,14 @@ def celu(x, alpha=1.0):
 
 def glu(x, axis=-1):
     x = convert_to_tensor(x)
+    dtype = x.dtype
     if x.shape[axis] % 2 != 0:
         raise ValueError(
             "axis size must be divisible by 2. "
             f"Received: x.shape={x.shape} with axis={axis}"
         )
     x1, x2 = np.split(x, 2, axis)
-    return x1 * (1 / (1 + np.exp(-x2)))
+    return (x1 * sigmoid(x2)).astype(dtype)
 
 
 def hard_tanh(x):

keras/src/ops/nn.py

@@ -704,7 +704,15 @@ def call(self, x):
         return backend.nn.glu(x, axis=self.axis)
 
     def compute_output_spec(self, x):
-        return KerasTensor(x.shape, dtype=x.dtype)
+        output_shape = list(x.shape)
+        if output_shape[self.axis] is not None:
+            if output_shape[self.axis] % 2 != 0:
+                raise ValueError(
+                    "axis size must be divisible by 2. "
+                    f"Received: x.shape={x.shape} with axis={self.axis}"
+                )
+            output_shape[self.axis] = output_shape[self.axis] // 2
+        return KerasTensor(output_shape, dtype=x.dtype)
 
 
 @keras_export(["keras.ops.glu", "keras.ops.nn.glu"])

keras/src/ops/nn_test.py

@@ -149,8 +149,8 @@ def test_celu(self):
         self.assertEqual(knn.celu(x).shape, (None, 2, 3))
 
     def test_glu(self):
-        x = KerasTensor([None, 2, 3])
-        self.assertEqual(knn.glu(x).shape, (None, 2, 3))
+        x = KerasTensor([None, 2, 4])
+        self.assertEqual(knn.glu(x).shape, (None, 2, 2))
 
     def test_tanh_shrink(self):
         x = KerasTensor([None, 2, 3])
@@ -851,8 +851,8 @@ def test_celu(self):
         self.assertEqual(knn.celu(x).shape, (1, 2, 3))
 
     def test_glu(self):
-        x = KerasTensor([1, 2, 3])
-        self.assertEqual(knn.glu(x).shape, (1, 2, 3))
+        x = KerasTensor([1, 2, 4])
+        self.assertEqual(knn.glu(x).shape, (1, 2, 2))
 
     def test_tanh_shrink(self):
         x = KerasTensor([1, 2, 3])
@@ -2734,9 +2734,6 @@ def test_glu(self, dtype):
         import jax.nn as jnn
         import jax.numpy as jnp
 
-        if dtype == "bfloat16":
-            self.skipTest("Weirdness with numpy")
-
         x = knp.ones((2), dtype=dtype)
         x_jax = jnp.ones((2), dtype=dtype)
         expected_dtype = standardize_dtype(jnn.glu(x_jax).dtype)