code: def rechunk_tensor(tensor, chunk_size):
      # TODO(hks): Provide a unify rechunk logic with mmap.
      cur_chunks = []
  
      out_nchunks = tensor.shape[0] // chunk_size
      row_nsplits = [chunk_size] * out_nchunks
      rest = tensor.shape[0] % chunk_size
      if rest >= out_nchunks:
          row_nsplits.append(rest)
      else:
          for i in range(tensor.shape[0] % chunk_size):
              row_nsplits[-i - 1] += 1
  
      tensor_cumnrows = np.cumsum([0] + list(tensor.nsplits[0]))
      offset = 0
      out_groups = []
      for split in row_nsplits:
          start_chunk_index = int(tensor_cumnrows.searchsorted(offset))
          start_chunk_index = (
              start_chunk_index - 1 if start_chunk_index != 0 else 0
          )
          end_chunk_index = int(
              tensor_cumnrows.searchsorted(offset + split) - 1
          )
          if start_chunk_index == end_chunk_index:
              t = tensor.chunks[start_chunk_index]
              slice_op = TensorSlice(
                  [
                      slice(
                          offset - tensor_cumnrows[start_chunk_index],
                          split
                          + offset
                          - tensor_cumnrows[end_chunk_index],
                      ),
                      slice(None),
                  ],
                  dtype=t.dtype,
              )
              out_groups.append(
                  [
                      slice_op.new_chunk(
                          [t],
                          shape=(split, t.shape[1]),
                          index=(len(cur_chunks), 0),
                          order=t.order,
                      )
                  ]
              )
          else:
              chunks = []
              start_chunk = tensor.chunks[start_chunk_index]
              start_slice = int(
                  offset - tensor_cumnrows[start_chunk_index]
              )
              slice_op = TensorSlice(
                  [slice(start_slice, None), slice(None)],
                  dtype=start_chunk.dtype,
              )
              chunks.append(
                  slice_op.new_chunk(
                      [start_chunk],
                      shape=(
                          start_chunk.shape[0] - start_slice,
                          start_chunk.shape[1],
                      ),
                      index=(0, 0),
                      order=start_chunk.order,
                  )
              )
              chunks.extend(
                  tensor.chunks[start_chunk_index + 1 : end_chunk_index]
              )
              end_chunk = tensor.chunks[end_chunk_index]
              end_slice = int(
                  split + offset - tensor_cumnrows[end_chunk_index]
              )
              slice_op_end = TensorSlice(
                  [slice(None, end_slice), slice(None)],
                  dtype=start_chunk.dtype,
              )
              chunks.append(
                  slice_op_end.new_chunk(
                      [end_chunk],
                      shape=(end_slice, end_chunk.shape[1]),
                      index=(end_chunk_index - start_chunk_index, 0),
                      order=end_chunk.order,
                  )
              )
              out_groups.append(chunks)
  
          offset += split
  
      return out_groups
  
file_path: mars/mars/learn/proxima/core.py
imports: import numpy
  from ...None import tensor
  from ...tensor.merge import TensorConcatenate
  from ...tensor.indexing import TensorSlice
label_desc: Use image node_id 1 for calling a global function with example usage: rechunk_tensor(tensor, chunk_size) and returns: out_groups