According to the document , the parameter axis should act as:
axis: Axis along which to sort. The default is -1 (the last axis). If None, the flattened array is used.
However, in the following code:
import keras
import numpy as np
x = np.random.normal((4, 4, 4))
print(keras.ops.argpartition(x, 0, axis=None))
The execution fails with the error message:
Traceback (most recent call last):
File "/mnt/data2/Users/sunshuo/2026/dltest_3.3/src/test.py", line 5, in <module>
print(keras.ops.argpartition(x, 0, axis=None))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/data2/Users/sunshuo/anaconda3/envs/dltest_to_raise_issue/lib/python3.12/site-packages/keras/src/ops/numpy.py", line 7827, in argpartition
return backend.numpy.argpartition(x, kth, axis)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/data2/Users/sunshuo/anaconda3/envs/dltest_to_raise_issue/lib/python3.12/site-packages/keras/src/backend/tensorflow/numpy.py", line 3306, in argpartition
x = swapaxes(x, axis, -1)
^^^^^^^^^^^^^^^^^^^^^
File "/mnt/data2/Users/sunshuo/anaconda3/envs/dltest_to_raise_issue/lib/python3.12/site-packages/keras/src/backend/tensorflow/numpy.py", line 2565, in swapaxes
axis1 = tf.where(axis1 < 0, tf.add(axis1, x_rank), axis1)
^^^^^^^^^
TypeError: '<' not supported between instances of 'NoneType' and 'int'
By the way, it works correctly with KerasTensor:
import keras
import numpy as np
x_sym = keras.layers.Input(shape=(4, 4))
print(keras.ops.argpartition(x_sym, 0, axis=None)) #<KerasTensor shape=(None, 4, 4), dtype=int32, sparse=False, ragged=False, name=keras_tensor_1>
According to the document , the parameter axis should act as:
However, in the following code:
The execution fails with the error message:
By the way, it works correctly with KerasTensor: