code: def _sort_init(
      ib,
      shape,
      axis,
      keys_in,
      keys_out,
      values_out=None,
      value_init_func=None,
  ):
      """Initialize the output buffers by copying from inputs"""
      axis_mul_before = 1
      axis_mul_after = 1
      if axis < 0:
          axis = len(shape) + axis
      for i, value in enumerate(shape, 0):
          if i < axis:
              axis_mul_before *= value
          elif i > axis:
              axis_mul_after *= value
  
      # Set up threading
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
      nthread_tx = max_threads
      nthread_bx = ceil_div(shape[axis], max_threads)
      nthread_by = axis_mul_before
      nthread_bz = axis_mul_after
  
      # Copy the keys_in to initial output
      with ib.new_scope():
          tx, bx, by, bz = _get_threads(
              ib, nthread_tx, nthread_bx, nthread_by, nthread_bz
          )
          tid = bx * nthread_tx + tx
          idx = (by * shape[axis] + tid) * axis_mul_after + bz
          with ib.if_scope(tid < shape[axis]):
              keys_out[idx] = keys_in[idx]
              if values_out is not None:
                  values_out[idx] = value_init_func(idx, tid)
  
      return axis_mul_before, axis_mul_after
  
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 3 for calling a global function with example usage: _sort_init(ib, shape, axis, keys_in, keys_out, values_out, value_init_func) and returns: axis_mul_before,  axis_mul_after