code: def _splitrange(n, k):
      """Calculates array lens of np.array_split().
  
      This is the equivalent of
      `[len(x) for x in np.array_split(range(n), k)]`.
      """
      base = n // k
      output = [base] * k
      rem = n - sum(output)
      for i in range(len(output)):
          if rem > 0:
              output[i] += 1
              rem -= 1
      assert rem == 0, (rem, output, n, k)
      assert sum(output) == n, (output, n, k)
      return output
  
imports: import itertools
  from abc import abstractmethod
  from enum import Enum
  from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar, Union
  from ray.data._internal.block_batching.block_batching import batch_blocks
  from ray.data._internal.execution.interfaces.task_context import TaskContext
  from ray.data._internal.output_buffer import BlockOutputBuffer
  from ray.data.block import Block, BlockAccessor, DataBatch
label_desc: Use image node_id 2 for calling a global function with example usage: _splitrange(n, k) and returns: output