docstr: Approach here would be almost like PSRSSorter, but there are definitely some differences
  Main processes are listed below:
  Stage 1, local sort and regular samples collected
  State 2, gather and merge samples, choose and broadcast p-1 pivots
  Stage 3, Local data is partitioned
  Stage 4: all *ith* classes are gathered and merged, sizes should be calculated as well
  Stage 5: collect sizes from partitions, calculate how to partition given kth
  Stage 6: partition on each partitions
  Stage 7: align if axis is given, and more than 1 dimension
docstr_tok: ["Approach","here","would","be","almost","like","PSRSSorter",",","but","there","are","definitely","some","differences","Main","processes","are","listed","below",":","Stage","1",",","local","sort","and","regular","samples","collected","State","2",",","gather","and","merge","samples",",","choose","and","broadcast","p-1","pivots","Stage","3",",","Local","data","is","partitioned","Stage","4",":","all","*","ith","*","classes","are","gathered","and","merged",",","sizes","should","be","calculated","as","well","Stage","5",":","collect","sizes","from","partitions",",","calculate","how","to","partition","given","kth","Stage","6",":","partition","on","each","partitions","Stage","7",":","align","if","axis","is","given",",","and","more","than","1","dimension"]
code: def _tile_psrs(cls, op, kth):
      """
      Approach here would be almost like PSRSSorter, but there are definitely some differences
      Main processes are listed below:
      Stage 1, local sort and regular samples collected
      State 2, gather and merge samples, choose and broadcast p-1 pivots
      Stage 3, Local data is partitioned
      Stage 4: all *ith* classes are gathered and merged, sizes should be calculated as well
      Stage 5: collect sizes from partitions, calculate how to partition given kth
      Stage 6: partition on each partitions
      Stage 7: align if axis is given, and more than 1 dimension
      """
      out_tensor = op.outputs[0]
      return_value, return_indices = op.return_value, op.return_indices
      # preprocess, to make sure chunk shape on axis are approximately same
      (
          in_tensor,
          axis_chunk_shape,
          out_idxes,
          need_align,
      ) = yield from cls.preprocess(op)
      axis_offsets = [0] + np.cumsum(
          in_tensor.nsplits[op.axis]
      ).tolist()[:-1]
  
      out_chunks, out_indices_chunks = [], []
      for out_idx in out_idxes:
          # stage 1: local sort and regular samples collected
          (
              sorted_chunks,
              indices_chunks,
              sampled_chunks,
          ) = cls.local_sort_and_regular_sample(
              op, in_tensor, axis_chunk_shape, axis_offsets, out_idx
          )
  
          # stage 2: gather and merge samples, choose and broadcast p-1 pivots
          concat_pivot_chunk = cls.concat_and_pivot(
              op,
              axis_chunk_shape,
              out_idx,
              sorted_chunks,
              sampled_chunks,
          )
  
          # stage 3: Local data is partitioned
          partition_chunks = cls.partition_local_data(
              op,
              axis_chunk_shape,
              sorted_chunks,
              indices_chunks,
              concat_pivot_chunk,
          )
  
          proxy_chunk = TensorShuffleProxy(
              dtype=partition_chunks[0].dtype
          ).new_chunk(partition_chunks, shape=())
  
          # stage 4: all *ith* classes are gathered and merged,
          # note that we don't need sort here, op.psrs_kinds[2] is None
          # force need_align=True to get sort info
          (
              partition_merged_chunks,
              partition_indices_chunks,
              sort_info_chunks,
          ) = cls.partition_merge_data(
              op, True, True, partition_chunks, proxy_chunk
          )
  
          # stage5, collect sort infos and calculate partition info for each partitions
          partition_info_chunks = cls.calc_paritions_info(
              op, kth, in_tensor.shape[op.axis], sort_info_chunks
          )
  
          # Stage 6: partition on each partitions
          (
              partitioned_chunks,
              partitioned_indices_chunks,
          ) = cls.partition_on_merged(
              op,
              need_align,
              partition_merged_chunks,
              partition_indices_chunks,
              partition_info_chunks,
          )
  
          if not need_align:
              if return_value:
                  out_chunks.extend(partitioned_chunks)
              if return_indices:
                  out_indices_chunks.extend(partitioned_indices_chunks)
          else:
              (
                  align_reduce_chunks,
                  align_reduce_indices_chunks,
              ) = cls.align_partitions_data(
                  op,
                  out_idx,
                  in_tensor,
                  partitioned_chunks,
                  partitioned_indices_chunks,
                  sort_info_chunks,
              )
              if return_value:
                  out_chunks.extend(align_reduce_chunks)
              if return_indices:
                  out_indices_chunks.extend(align_reduce_indices_chunks)
  
      new_op = op.copy()
      nsplits = list(in_tensor.nsplits)
      if not need_align:
          nsplits[op.axis] = (np.nan,) * axis_chunk_shape
      kws = []
      if return_value:
          kws.append(
              {
                  "shape": out_tensor.shape,
                  "order": out_tensor.order,
                  "chunks": out_chunks,
                  "nsplits": tuple(nsplits),
                  "dtype": out_tensor.dtype,
                  "type": "partitioned",
              }
          )
      if return_indices:
          kws.append(
              {
                  "shape": out_tensor.shape,
                  "order": TensorOrder.C_ORDER,
                  "chunks": out_indices_chunks,
                  "nsplits": tuple(nsplits),
                  "dtype": np.dtype(np.int64),
                  "type": "argpartition",
              }
          )
      return new_op.new_tensors(op.inputs, kws=kws)
  
imports: import itertools
  import numpy
  from ...None import opcodes
  from ...core import recursive_tile
  from ...serialization.serializables import FieldTypes, KeyField, Int32Field, StringField, ListField, BoolField, AnyField
  from ...utils import has_unknown_shape, flatten, stack_back
  from ...core import ENTITY_TYPE, ExecutableTuple
  from ..operands import TensorOperand, TensorShuffleProxy
  from ..core import TENSOR_TYPE, TENSOR_CHUNK_TYPE, TensorOrder
  from ..array_utils import as_same_device, device
  from ..datasource import tensor
  from ..utils import validate_axis, validate_order
  from .psrs import TensorPSRSOperandMixin
label_desc: Use image node_id 14 for calling the TensorPartition obj's underlying member method code with example usage: obj._tile_psrs(cls, op, kth) and returns: new_op