code: def _odd_even_sort(
      ib,
      size,
      axis_mul_before,
      axis_mul_after,
      is_ascend,
      keys,
      keys_swap,
      values=None,
      values_swap=None,
  ):
      nthread_tx = block_size // 2
      nthread_bx = ceil_div(size, block_size)
      nthread_by = axis_mul_before
      nthread_bz = axis_mul_after
      with ib.new_scope():
          ib.scope_attr(tvm.tir.const(0), "hand_threaded", 0)
          tx, bx, by, bz = _get_threads(
              ib, nthread_tx, nthread_bx, nthread_by, nthread_bz
          )
          tid = 2 * tx
          start = bx * block_size
  
          ## Create shared memory as syncable thread scratch space
          tmp_keys_swap = ib.allocate(
              keys_swap.dtype,
              (block_size,),
              name="temp_keys_swap",
              scope="shared",
          )
          if values_swap is not None:
              tmp_values_swap = ib.allocate(
                  values_swap.dtype,
                  (block_size,),
                  name="temp_values_swap",
                  scope="shared",
              )
  
          ## Create thread local data for swapping
          temp_keys = ib.allocate(
              keys_swap.dtype, (1,), name="temp_keys", scope="local"
          )
          if values_swap is not None:
              temp_values = ib.allocate(
                  values_swap.dtype,
                  (1,),
                  name="temp_values",
                  scope="local",
              )
  
          temp_cond1 = ib.allocate(
              keys_swap.dtype, (1,), name="temp_cond1", scope="local"
          )
          temp_cond2 = ib.allocate(
              keys_swap.dtype, (1,), name="temp_cond2", scope="local"
          )
          # Copy data to scratch space
          base_idx = by * size * axis_mul_after + bz
          with ib.for_range(0, 2) as n:
              with ib.if_scope((tid + n + start) < size):
                  tmp_keys_swap[tid + n] = keys[
                      base_idx + (tid + n + start) * axis_mul_after
                  ]
                  if values_swap is not None:
                      tmp_values_swap[tid + n] = values[
                          base_idx + (tid + n + start) * axis_mul_after
                      ]
  
          ib.emit(
              tvm.tir.Call(
                  None,
                  "tir.tvm_storage_sync",
                  tvm.runtime.convert(["shared"]),
              )
          )
  
          idxm = tvm.tir.indexmod
          # OddEvenTransposeSort
          current_sort_num = tvm.tir.min(block_size, size - start)
          with ib.for_range(0, current_sort_num) as k:
              n = idxm(tid + k, 2)
              with ib.if_scope(tid + n < current_sort_num - 1):
                  temp_cond1[0] = tmp_keys_swap[tid + n]
                  temp_cond2[0] = tmp_keys_swap[tid + n + 1]
                  if is_ascend:
                      cond = temp_cond1[0] > temp_cond2[0]
                  else:
                      cond = temp_cond1[0] < temp_cond2[0]
                  with ib.if_scope(cond):
                      temp_keys[0] = tmp_keys_swap[tid + n]
                      tmp_keys_swap[tid + n] = tmp_keys_swap[
                          tid + n + 1
                      ]
                      tmp_keys_swap[tid + n + 1] = temp_keys[0]
                      if values_swap is not None:
                          temp_values[0] = tmp_values_swap[tid + n]
                          tmp_values_swap[tid + n] = tmp_values_swap[
                              tid + n + 1
                          ]
                          tmp_values_swap[tid + n + 1] = temp_values[0]
              ib.emit(
                  tvm.tir.Call(
                      None,
                      "tir.tvm_storage_sync",
                      tvm.runtime.convert(["shared"]),
                  )
              )
  
          ## Copy sorted data to output
          with ib.for_range(0, 2) as n:
              with ib.if_scope(tid + n + start < size):
                  keys[
                      base_idx + (tid + n + start) * axis_mul_after
                  ] = tmp_keys_swap[tid + n]
                  keys_swap[
                      base_idx + (tid + n + start) * axis_mul_after
                  ] = tmp_keys_swap[tid + n]
                  if values_swap is not None:
                      values[
                          base_idx + (tid + n + start) * axis_mul_after
                      ] = tmp_values_swap[tid + n]
                      values_swap[
                          base_idx + (tid + n + start) * axis_mul_after
                      ] = tmp_values_swap[tid + n]
  
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 4 for calling a global function with example usage: _odd_even_sort(ib, size, axis_mul_before, axis_mul_after, is_ascend, keys, keys_swap, values, values_swap) without return types