def load_array(self, pattern: str):
    """Load sharded array as if it was loaded from multiple processes."""
    process_count = jax.process_count()
    arrays = []
    for process_index in range(process_count):
      arrays.append(
          data_utils.load_sharded_array(pattern,
                                        process_count * self.memory_reduction,
                                        process_index))
    array = np.stack(arrays, axis=0)
    shape = (-1,) + arrays[0].shape[1:]
    array = array.reshape(shape)
    return array
Load sharded array as if it was loaded from multiple processes.
language/mentionmemory/encoders/mention_memory_encoder.py