`keras.ops.repeat` cannot return an exptected shape when `x` is a `KerasTensor` and the `axis` is `None`

[Zhaopudark]

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), 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：

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:

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:

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.

[pandyah5]

I can look into this and report my findings in a few hours

[Grvzard]

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.

[pandyah5]

Edited: [Was confused about the expected dimensions of the output but I found the mistake in my logic]