code: def _topk_helper(
      xp,
      a,
      k,
      axis=-1,
      largest=True,
      sorted=True,
      order=None,
      indices=None,
      axis_offset=None,
      return_value=True,
      return_indices=False,
  ):
      size = a.shape[axis]
      base_slc = (slice(None),) * axis
      kw = {}
      if order is not None:
          kw["order"] = order
  
      ap = None
      if return_indices:
          # do partition
          if largest:
              if k < size:
                  length = size - k
                  ap = xp.argpartition(a, length, axis=axis, **kw)[
                      base_slc + (slice(-k, None),)
                  ]
                  av = xp.take_along_axis(a, ap, axis)
                  if indices is not None:
                      ap = xp.take_along_axis(indices, ap, axis)
              else:
                  av = a
                  if indices is not None:
                      ap = indices
                  else:
                      ap = _gen_indices(a.shape, axis, xp)
              if sorted:
                  # sort then reverse
                  argsort_result = xp.argsort(av, axis=axis, **kw)[
                      base_slc + (slice(None, None, -1),)
                  ]
                  ap = xp.take_along_axis(ap, argsort_result, axis)
                  av = xp.take_along_axis(av, argsort_result, axis)
          else:
              if k < size:
                  ap = xp.argpartition(a, k, axis=axis, **kw)[
                      base_slc + (slice(k),)
                  ]
                  av = xp.take_along_axis(a, ap, axis)
                  if indices is not None:
                      ap = xp.take_along_axis(indices, ap, axis)
              else:
                  av = a
                  if indices is not None:
                      ap = indices
                  else:
                      ap = _gen_indices(a.shape, axis, xp)
              if sorted:
                  argsort_result = xp.argsort(av, axis=axis, **kw)
                  ap = xp.take_along_axis(ap, argsort_result, axis)
                  av = xp.take_along_axis(av, argsort_result, axis)
          if axis_offset:
              ap = ap + axis_offset
      else:
          assert return_value
          if largest:
              if k < size:
                  length = size - k
                  av = xp.partition(a, length, axis=axis, **kw)[
                      base_slc + (slice(-k, None),)
                  ]
              else:
                  av = a
              if sorted:
                  # sort then reverse
                  av = xp.sort(av, axis=axis, **kw)[
                      base_slc + (slice(None, None, -1),)
                  ]
          else:
              if k < size:
                  av = xp.partition(a, k, axis=axis, **kw)[
                      base_slc + (slice(k),)
                  ]
              else:
                  av = a
              if sorted:
                  av = xp.sort(av, axis=axis, **kw)
  
      return av, ap
  
imports: import itertools
  import numpy
  from ...None import opcodes
  from ...config import options
  from ...core import ExecutableTuple, recursive_tile
  from ...core.operand import OperandStage
  from ...serialization.serializables import FieldTypes, KeyField, Int64Field, Int32Field, BoolField, StringField, ListField
  from ...utils import ceildiv, flatten
  from ..array_utils import as_same_device, device
  from ..core import TensorOrder
  from ..datasource import tensor
  from ..operands import TensorOperand, TensorOperandMixin
  from ..utils import validate_axis, validate_order
  from .sort import _validate_sort_psrs_kinds
label_desc: Use image node_id 2 for calling a global function with example usage: _topk_helper(xp, a, k, axis, largest, sorted, order, indices, axis_offset, return_value, return_indices) and returns: av,  ap