code: def sort_by_key_ir(
      keys_in,
      values_in,
      keys_out,
      values_out,
      keys_out_swap,
      values_out_swap,
      axis,
      is_ascend,
  ):
      """Low level IR to do sort by key on the GPU.
  
      Parameters
      ----------
      keys_in: Buffer
          Buffer of input keys.
  
      values_in: Buffer
          Buffer of input keys.
  
      keys_out : Buffer
          Buffer of output sorted keys.
  
      values_out : Buffer
          Buffer of output sorted values.
  
      keys_out_swap : Buffer
          Output buffer of values with same shape as keys_in to use as swap.
  
      values_out_swap : Buffer
          Output buffer of values with same shape as values_in to use as swap.
  
      axis : Int
          Axis long which to sort the input tensor.
  
      is_ascend : Boolean
          Whether to sort in ascending or descending order.
  
      indicess_out : Buffer
          Output buffer of indices of sorted tensor with same shape as keys_in.
  
      values_out_swap : Buffer
          Output buffer of indices with same shape as keys_in to use as swap.
  
      Returns
      -------
      stmt : Stmt
          The result IR statement.
      """
      ib = tvm.tir.ir_builder.create()
      shape = keys_in.shape
  
      keys_in = ib.buffer_ptr(keys_in)
      values_in = ib.buffer_ptr(values_in)
      keys_out = ib.buffer_ptr(keys_out)
      keys_out_swap = ib.buffer_ptr(keys_out_swap)
      values_out = ib.buffer_ptr(values_out)
      values_out_swap = ib.buffer_ptr(values_out_swap)
  
      with ib.if_scope(shape[axis] > 0):
          axis_mul_before, axis_mul_after = _sort_init(
              ib,
              shape,
              axis,
              keys_in,
              keys_out,
              values_out,
              value_init_func=lambda idx, _: values_in[idx],
          )
  
          _sort_common(
              ib,
              shape[axis],
              axis_mul_before,
              axis_mul_after,
              is_ascend,
              keys_out,
              keys_out_swap,
              values=values_out,
              values_swap=values_out_swap,
          )
      return ib.get()
  
imports: import tvm
  from tvm import te
  from .injective import schedule_injective_from_existing
  from ..transform import strided_slice, transpose
  from ..None import tag
  from ..utils import ceil_div, swap
  from ..math import cast, ceil_log2
label_desc: Use image node_id 7 for calling a global function with example usage: sort_by_key_ir(keys_in, values_in, keys_out, values_out, keys_out_swap, values_out_swap, axis, is_ascend) and returns: ib