def load_sharded_array(
    pattern: Union[str, Sequence[str]],
    stride: int,
    offset: int,
) -> np.ndarray:
  """Load and concatenate numpy arrays according to pattern.

  Args:
    pattern: pattern of array filenames to be loaded.
    stride: stride to traverse paths to be loaded.
    offset: starting offset into paths to be loaded.

  Returns:
    Loaded and concatenated array.
  """
  paths = tf.io.gfile.glob(pattern)
  array_list = []
  for path in paths[offset::stride]:
    logging.info('Loading %s on to process %d', path, jax.process_index())
    with tf.io.gfile.GFile(path, 'rb') as f:
      array_list.append(np.load(f))

  array = np.concatenate(array_list, axis=0)
  return array
Load and concatenate numpy arrays according to pattern.

  Args:
    pattern: pattern of array filenames to be loaded.
    stride: stride to traverse paths to be loaded.
    offset: starting offset into paths to be loaded.

  Returns:
    Loaded and concatenated array.
  
language/mentionmemory/utils/data_utils.py