def np_to_file(list_of_nparrays, file_path, compresslevel):
  """Save numpy arrays to file_path with gzipping and failure protection."""
  # Pickle to tmp file and overwrite to prevent writing partial files.
  tmp_file_path = file_path + '._tmp_'
  with tf.io.gfile.GFile(tmp_file_path, 'wb') as f:
    with gzip.GzipFile(fileobj=f, compresslevel=compresslevel) as gzipf:
      for x in list_of_nparrays:
        np.save(gzipf, x, allow_pickle=False)
  # Moving a file is much less error-prone than pickling large files.
  tf.io.gfile.rename(tmp_file_path, file_path, overwrite=True)
Save numpy arrays to file_path with gzipping and failure protection.
trax/layers/base.py