def _sort(*operands: TfVal, dimension: int, is_stable: bool, num_keys: int) -> Tuple[(TfVal, ...)]:
    if (not _enable_xla):
        raise _xla_disabled_error('sort')
    assert (1 <= num_keys <= len(operands))
    assert (0 <= dimension < len(operands[0].shape)), f'Invalid {dimension} for ndim {len(operands[0].shape)}'

    def lexicographic_comparator_old(*tf_args: TfVal) -> TfVal:
        assert (len(tf_args) == (2 * len(operands)))
        inside_comparison = None
        for key_idx in range((num_keys - 1), (- 1), (- 1)):
            a = tf_args[(2 * key_idx)]
            b = tf_args[((2 * key_idx) + 1)]
            a_lt_b = tf.math.less(a, b)
            if (inside_comparison is None):
                inside_comparison = a_lt_b
            else:
                inside_comparison = tf.math.logical_or(a_lt_b, tf.math.logical_and(tf.math.equal(a, b), inside_comparison))
        return inside_comparison
    comparator_spec: List[tf.TensorSpec] = []
    comparator_jax_in_avals: List[core.AbstractValue] = []
    for op in operands:
        o_spec = tf.TensorSpec((), dtype=op.dtype)
        comparator_spec.extend([o_spec, o_spec])
        o_aval = core.ShapedArray((), to_jax_dtype(op.dtype))
        comparator_jax_in_avals.extend([o_aval, o_aval])

    def lexicographic_comparator(*tf_args: TfVal) -> TfVal:
        return _convert_jax_impl(lax._sort_lt_comparator, multiple_results=False)(*tf_args, _in_avals=comparator_jax_in_avals, _out_aval=core.ShapedArray((), np.bool_), num_keys=num_keys)
    xla_comparator_computation = tf.function(lexicographic_comparator, autograph=False).get_concrete_function(*comparator_spec)
    results = tfxla.variadic_sort(operands, dimension=dimension, is_stable=is_stable, comparator=xla_comparator_computation)
    return results
def _sort(*operands: TfVal, dimension: int, is_stable: bool, num_keys: int) -> Tuple[(TfVal, ...)]:
    if (not _enable_xla):
        raise _xla_disabled_error('sort')
    assert (1 <= num_keys <= len(operands))
    assert (0 <= dimension < len(operands[0].shape)), f'Invalid {dimension} for ndim {len(operands[0].shape)}'

    def lexicographic_comparator_old(*tf_args: TfVal) -> TfVal:
        assert (len(tf_args) == (2 * len(operands)))
        inside_comparison = None
        for key_idx in range((num_keys - 1), (- 1), (- 1)):
            a = tf_args[(2 * key_idx)]
            b = tf_args[((2 * key_idx) + 1)]
            a_lt_b = tf.math.less(a, b)
            if (inside_comparison is None):
                inside_comparison = a_lt_b
            else:
                inside_comparison = tf.math.logical_or(a_lt_b, tf.math.logical_and(tf.math.equal(a, b), inside_comparison))
        return inside_comparison
    comparator_spec: List[tf.TensorSpec] = []
    comparator_jax_in_avals: List[core.AbstractValue] = []
    for op in operands:
        o_spec = tf.TensorSpec((), dtype=op.dtype)
        comparator_spec.extend([o_spec, o_spec])
        o_aval = core.ShapedArray((), to_jax_dtype(op.dtype))
        comparator_jax_in_avals.extend([o_aval, o_aval])

    def lexicographic_comparator(*tf_args: TfVal) -> TfVal:
        return _convert_jax_impl(lax._sort_lt_comparator, multiple_results=False)(*tf_args, _in_avals=comparator_jax_in_avals, _out_aval=core.ShapedArray((), np.bool_), num_keys=num_keys)
    xla_comparator_computation = tf.function(lexicographic_comparator, autograph=False).get_concrete_function(*comparator_spec)
    results = tfxla.variadic_sort(operands, dimension=dimension, is_stable=is_stable, comparator=xla_comparator_computation)
    return results