code: def sort_by_key(keys, values, axis=-1, is_ascend=1):
      """Sort values with respect to keys. Both keys and values will
       be sorted and returned.
  
      Parameters
      ----------
      keys: tvm.te.Tensor
          The input keys.
  
      values : tvm.te.Tensor,
          The input values.
  
      axis : int, optional
          Axis long which to sort the input tensor.
  
      is_ascend : boolean, optional
          Whether to sort in ascending or descending order.
  
      Returns
      -------
      keys_sorted : tvm.te.Tensor
          The sorted keys
  
      values_sorted : tvm.te.Tensor
          The values sorted with respect to the keys
      """
      keys_buf = tvm.tir.decl_buffer(
          keys.shape, keys.dtype, "keys_buf", data_alignment=8
      )
      values_buf = tvm.tir.decl_buffer(
          values.shape, values.dtype, "values_buf", data_alignment=8
      )
  
      out_bufs = [
          tvm.tir.decl_buffer(
              keys.shape, keys.dtype, "keys_buf", data_alignment=8
          ),
          tvm.tir.decl_buffer(
              values.shape, values.dtype, "values_buf", data_alignment=8
          ),
          tvm.tir.decl_buffer(
              keys.shape, keys.dtype, "keys_swap_buf", data_alignment=8
          ),
          tvm.tir.decl_buffer(
              values.shape,
              values.dtype,
              "values_swap_buf",
              data_alignment=8,
          ),
      ]
      out = te.extern(
          [keys.shape, values.shape, keys.shape, values.shape],
          [keys, values],
          lambda ins, outs: sort_by_key_ir(
              ins[0],
              ins[1],
              outs[0],
              outs[1],
              outs[2],
              outs[3],
              axis,
              is_ascend,
          ),
          in_buffers=[keys_buf, values_buf],
          out_buffers=out_bufs,
          dtype=[keys.dtype, values.dtype],
          name="sort_by_key",
          tag="sort_by_key",
      )
      return out[0], out[1]
  
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 17 for calling a global function with example usage: sort_by_key(keys, values, axis, is_ascend) and returns: out,  out