def save_sharded_array(array: np.ndarray, prefix: str, num_shards: int,
                       stride: int, offset: int, shard_size_divisible: int):
  """Save numpy array into multiple shards.

  Input array will be divided and saved into files with names
  "prefix-X-of-num_shards", where `X` is in [offset, offset + stride, ...].

  Args:
    array: numpy array to be saved.
    prefix: output filenames prefix.
    num_shards: total number of shards.
    stride: stride for which shards save to the input array.
    offset: starting offset for which shards save to the input array.
    shard_size_divisible: ensure that the first dimension of each shard is
      divisible by this number via zero padding at the end.
  """
  if num_shards % stride != 0:
    raise ValueError('`num_shards` must to be divisible by the stride.')

  actual_num_shards = num_shards // stride

  actual_data_per_shard = np.full((actual_num_shards),
                                  array.shape[0] // actual_num_shards)
  actual_data_per_shard[:array.shape[0] % actual_num_shards] += 1

  shard_size = (
      math.ceil(actual_data_per_shard.max() / shard_size_divisible) *
      shard_size_divisible)

  array_shard = np.zeros(
      [shard_size] + list(array.shape[1:]), dtype=array.dtype)

  start_index = 0
  for i, shard_index in enumerate(range(offset, num_shards, stride)):
    end_index = min(array.shape[0], start_index + actual_data_per_shard[i])
    actual_size = end_index - start_index
    output_path = '%s-%05d-of-%05d' % (prefix, shard_index, num_shards)

    array_shard[:] = 0
    array_shard[:actual_size] = array[start_index:end_index]

    with tf.io.gfile.GFile(output_path, 'wb') as f:
      np.save(f, array_shard)

    start_index = end_index
Save numpy array into multiple shards.

  Input array will be divided and saved into files with names
  "prefix-X-of-num_shards", where `X` is in [offset, offset + stride, ...].

  Args:
    array: numpy array to be saved.
    prefix: output filenames prefix.
    num_shards: total number of shards.
    stride: stride for which shards save to the input array.
    offset: starting offset for which shards save to the input array.
    shard_size_divisible: ensure that the first dimension of each shard is
      divisible by this number via zero padding at the end.
  
language/mentionmemory/utils/data_utils.py