keras/src/backend/torch/numpy.py/argsort
def argsort(x, axis=-1):
    x = convert_to_tensor(x)

    # TODO: torch.argsort doesn't support bool
    if standardize_dtype(x.dtype) == "bool":
        x = cast(x, "uint8")

    if axis is None:
        axis = -1
        x = x.reshape(-1)
    return cast(torch.argsort(x, dim=axis, stable=True), dtype="int32")