#5709·burn

pytorch-reader: read a checkpoint from bytes, not only from a path

Author: antimoraCreated Sep 16, 2026Updated Sep 17, 2026
Labelsenhancementstore

Summary

PytorchReader opens checkpoints by path only. from_reader exists but is documented, and implemented, as plain-pickle only:

The reader must hold a plain pickle: tensor data lives outside the pickle in every PyTorch container, so a checkpoint with tensors must be loaded from a file.

The second half does not follow from the first. Tensor bytes live outside the pickle, inside the container, and a container is just bytes: a ZIP, a legacy stream or a TAR held in a Vec<u8> has everything the file has. Nothing in the parsing needs a path.

Where it bites

  • A checkpoint fetched over HTTP, pulled out of an archive, or decrypted, and handed over as bytes, has to be written to a temp file first just to be read.
  • The reader's own structure already allows it: a Container enum with File(PathBuf) and Memory(Arc<Vec<u8>>) variants, whose open() returns a Box<dyn Read + Seek> over either, is enough for ZIP and legacy, and I have run the parser that way.
  • The other stores in burn-store (SafetensorsStore::from_bytes, BurnpackStore::from_bytes) take bytes; PytorchStore is the odd one out.

Proposed shape

rust
impl PytorchReader {
    pub fn from_bytes(bytes: Vec<u8>, top_level_key: Option<&str>) -> Result<Self>;
}

Detection already reads only a header; ZipSource can hold ZipArchive<Cursor<Arc<Vec<u8>>>> as easily as ZipArchive<BufReader<File>>, legacy reads seek within the buffer, TAR is already read into memory at open. The Arc matters: every tensor's read closure shares the buffer instead of cloning it. from_reader can then either stay as the plain-pickle entry point with its doc fixed, or read the stream to the end and delegate.

Laziness is unchanged: Tensor::read still gathers on demand, from the buffer instead of the file.

Also

Whatever happens to the constructor, the from_reader doc should say "this function does not read containers" rather than "containers cannot be read from memory".