Repository: keras-team/keras

Issue:
 `keras.ops.repeat` cannot return an exptected shape when `x` is a `KerasTensor` and the `axis` is `None`
Hello. Thank you for your contributions and maintenance for the best Keras.

I'm following the instructions of [Conditional GAN (code samples, uses Keras 3)](https://keras.io/examples/generative/conditional_gan/), and focusing on the `keras.ops.repeat` function that is used in it.

I have found, maybe, if the input tensor of  `keras.ops.repeat`  is a symbolic tensor, i.e., the `keras.KerasTensor`, and the arg `axis` is `None`, the returned one will not be my expected one.

As the following：
```python
batch_size = 64
class_num = 10
a = keras.KerasTensor(shape=(batch_size, class_num), dtype=tf.float32)
a = a[:, :, None, None] # [B,10,1,1]
b = keras.ops.repeat(a, repeats=[28 * 28])
print(b.shape)# (784,)
# expected output: (501760,)
```

If assign `axis`, it works as expected:
```python
a = keras.KerasTensor(shape=(batch_size, class_num), dtype=tf.float32)
a = a[:, :, None, None] # [B,10,1,1]
b = keras.ops.repeat(a, repeats=[28 * 28],axis=0)
print(b.shape)# (784, 10, 1, 1)
# expected output: (784, 10, 1, 1)
```

If not use the symbolic tensor, it also works as expected:
```python
a = keras.random.normal(shape=(batch_size, class_num), dtype=tf.float32)
a = a[:, :, None, None] # [B,10,1,1]
b = keras.ops.repeat(a, repeats=[28 * 28])
print(b.shape)# (501760,)
# expected output: (501760,)
```

So, is the above a bug? 

And my environment is:
- Keras: Version: 3.3.3
- Numpy: Version: 1.26.4
- TensorFlow: Version: 2.16.1

Thanks in advance.


Discussion:
I can look into this and report my findings in a few hours
This is due to an oversight caused by the different ways Keras and other backends handle the `repeats` parameter.
You can submit a PR after you solve it.
Edited: [Was confused about the expected dimensions of the output but I found the mistake in my logic]



Patch:
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
index e8289ee2ebd7..69dbe23076d3 100644
--- a/keras/src/ops/numpy.py
+++ b/keras/src/ops/numpy.py
@@ -4345,26 +4345,44 @@ def call(self, x):
 
     def compute_output_spec(self, x):
         x_shape = list(x.shape)
+        repeats = self.repeats
+        if isinstance(repeats, int):
+            repeats = [repeats]
+        repeats_size = len(repeats)
+        broadcast = repeats_size == 1
+
         if self.axis is None:
             if None in x_shape:
                 return KerasTensor([None], dtype=x.dtype)
 
             x_flatten_size = int(np.prod(x_shape))
-            if isinstance(self.repeats, int):
-                output_shape = [x_flatten_size * self.repeats]
+            if broadcast:
+                output_shape = [x_flatten_size * repeats[0]]
+            elif repeats_size != x_flatten_size:
+                raise ValueError(
+                    "Size of `repeats` and "
+                    "dimensions of `x` after flattening should be compatible. "
+                    f"Received: {repeats_size} and {x_flatten_size}"
+                )
             else:
-                output_shape = [int(np.sum(self.repeats))]
+                output_shape = [int(np.sum(repeats))]
             return KerasTensor(output_shape, dtype=x.dtype)
 
         size_on_ax = x_shape[self.axis]
+        if size_on_ax is None:
+            return KerasTensor(x_shape, dtype=x.dtype)
+
         output_shape = x_shape
-        if isinstance(self.repeats, int):
-            if size_on_ax is None:
-                output_shape[self.axis] = None
-            else:
-                output_shape[self.axis] = size_on_ax * self.repeats
+        if broadcast:
+            output_shape[self.axis] = size_on_ax * repeats[0]
+        elif size_on_ax != repeats_size:
+            raise ValueError(
+                "Size of `repeats` and "
+                f"dimensions of `axis {self.axis} of x` should be compatible. "
+                f"Received: {repeats_size} and {x_shape}"
+            )
         else:
-            output_shape[self.axis] = int(np.sum(self.repeats))
+            output_shape[self.axis] = int(np.sum(repeats))
         return KerasTensor(output_shape, dtype=x.dtype)
 
 


Test patch:
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
index e257cde4af5b..ab730608fc4f 100644
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -1364,7 +1364,7 @@ def test_repeat(self):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.repeat(x, 2).shape, (None,))
         self.assertEqual(knp.repeat(x, 3, axis=1).shape, (None, 9))
-        self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (3, 3))
+        self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (None, 3))
         self.assertEqual(knp.repeat(x, 2, axis=0).shape, (None, 3))
 
     def test_reshape(self):
@@ -1875,9 +1875,15 @@ def test_reciprocal(self):
     def test_repeat(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.repeat(x, 2).shape, (12,))
+        self.assertEqual(knp.repeat(x, [2]).shape, (12,))
         self.assertEqual(knp.repeat(x, 3, axis=1).shape, (2, 9))
         self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (3, 3))
 
+        with self.assertRaises(ValueError):
+            knp.repeat(x, [1, 1])
+        with self.assertRaises(ValueError):
+            knp.repeat(x, [1, 1, 1], axis=0)
+
     def test_reshape(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.reshape(x, (3, 2)).shape, (3, 2))
@@ -3902,6 +3908,10 @@ def test_reciprocal(self):
     def test_repeat(self):
         x = np.array([[1, 2], [3, 4]])
         self.assertAllClose(knp.repeat(x, 2), np.repeat(x, 2))
+        self.assertAllClose(
+            knp.Repeat(np.array([2]))(x),
+            np.repeat(x, np.array([2])),
+        )
         self.assertAllClose(knp.repeat(x, 3, axis=1), np.repeat(x, 3, axis=1))
         self.assertAllClose(
             knp.repeat(x, np.array([1, 2]), axis=-1),