def wrap_sort(*args): return xla.variadic_sort( args, # Pass the arguments as a tuple comparator=compare_lt, dimension=dimension, is_stable=True)
def wrap_sort(x): return xla.variadic_sort([x], dimension=dimension, is_stable=False, comparator=compare_lt)
def wrap_sort(*args): return xla.variadic_sort( args, # Pass the arguments as a tuple comparator=compare_lexicographic, dimension=0, is_stable=False)