keras/src/ops/numpy.py/Argsort/compute_output_spec
class Argsort:    def compute_output_spec(self, x):
        if self.axis is None:
            return KerasTensor([int(np.prod(x.shape))], dtype="int32")
        return KerasTensor(x.shape, dtype="int32")