burn_pack/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2
3//! # Burn Pack
4//!
5//! The **burnpack** binary serialization format for the Burn deep learning framework.
6//!
7//! `burn-pack` is intentionally minimal and tensor-library-agnostic: it depends only on
8//! `burn-std` (for [`DType`] / [`Bytes`]), `serde`, and a CBOR codec. It knows how to read
9//! and write the burnpack container format but has no notion of Burn modules or tensors.
10//! Higher layers (e.g. `burn-core`) bridge between [`Tensor`] entries and their own
11//! tensor/snapshot types.
12//!
13//! Write a pack with [`Writer`], read one with [`Reader`]; both operate on [`Tensor`]
14//! entries that carry the format-level metadata plus a lazy provider of the raw
15//! little-endian bytes.
16//!
17//! ```
18//! use burn_pack::{Bytes, DType, Reader, Tensor, Writer};
19//!
20//! // A 2x2 f32 tensor, as raw little-endian bytes.
21//! let raw: Vec<u8> = [1.0f32, 2.0, 3.0, 4.0]
22//!     .iter()
23//!     .flat_map(|v| v.to_le_bytes())
24//!     .collect();
25//! let tensor = Tensor::new(
26//!     "weight".to_string(),
27//!     DType::F32,
28//!     vec![2, 2],
29//!     Some(42), // optional param id
30//!     Bytes::from_bytes_vec(raw),
31//! );
32//!
33//! // Write to an in-memory buffer ...
34//! let packed = Writer::new(vec![tensor])
35//!     .with_metadata("producer", "burn-pack docs")
36//!     .into_bytes()
37//!     .unwrap();
38//!
39//! // ... and read it back.
40//! let reader = Reader::from_bytes(packed).unwrap();
41//! assert_eq!(reader.metadata()["producer"], "burn-pack docs");
42//! // Consume the reader to get the tensors (zero-copy views into the source).
43//! let tensors = reader.into_tensors().unwrap();
44//! assert_eq!(tensors.len(), 1);
45//! assert_eq!(tensors[0].name, "weight");
46//! assert_eq!(tensors[0].shape.to_vec(), vec![2, 2]);
47//! assert_eq!(tensors[0].param_id, Some(42));
48//! ```
49//!
50//! # File format
51//!
52//! A burnpack file has three parts: a fixed-size header, a CBOR metadata blob, and a
53//! 256-byte-aligned tensor data section. All multi-byte integers are little-endian.
54//!
55//! ```text
56//! ┌──────────────────────────────────────────────────────────────┐
57//! │ Header — 10 bytes ([`HEADER_SIZE`])                           │
58//! │   magic         : u32  — 0x4255524E "BURN" ([`MAGIC_NUMBER`]) │
59//! │   version       : u16  — format version ([`FORMAT_VERSION`])  │
60//! │   metadata_size : u32  — byte length of the CBOR metadata     │
61//! ├──────────────────────────────────────────────────────────────┤
62//! │ Metadata — CBOR, `metadata_size` bytes                       │
63//! │   tensors : map<name, descriptor>                            │
64//! │     dtype        : [`DType`]                                  │
65//! │     shape        : list<u64>                                  │
66//! │     data_offsets : (start, end)  relative to the data section │
67//! │     param_id     : optional u64  (training-state identity)    │
68//! │   metadata : map<string, string>  user key/value pairs       │
69//! ├──────────────────────────────────────────────────────────────┤
70//! │ Padding to the next 256-byte boundary                        │
71//! │   ([`aligned_data_section_start`])                            │
72//! ├──────────────────────────────────────────────────────────────┤
73//! │ Tensor data section                                          │
74//! │   each tensor's bytes start on a 256-byte boundary           │
75//! │   ([`TENSOR_ALIGNMENT`]) for aligned, lazy file-backed       │
76//! │   loading (see [`Bytes::from_file`]).                         │
77//! │   tensors sliced zero-copy.                                   │
78//! └──────────────────────────────────────────────────────────────┘
79//! ```
80//!
81//! ## Why 256-byte alignment
82//!
83//! Aligning every tensor to a 256-byte boundary ([`TENSOR_ALIGNMENT`]) lets a reader
84//! memory-map the file and hand out tensor slices without copying, while satisfying the
85//! alignment requirements of every element type (including 8-byte `f64`), cache lines,
86//! and GPU coalesced access. 256 bytes matches the choice made by GGUF, MLX, ncnn, and
87//! other major formats.
88//!
89//! ## Safety limits
90//!
91//! Reading is hardened against malicious or corrupt inputs. The reader rejects files
92//! that exceed any of the following before allocating for them:
93//!
94//! - [`MAX_METADATA_SIZE`] — largest CBOR metadata blob
95//! - [`MAX_TENSOR_COUNT`] — largest number of tensors
96//! - [`MAX_TENSOR_SIZE`] — largest single tensor
97//! - [`MAX_CBOR_RECURSION_DEPTH`] — deepest CBOR nesting (stack-overflow guard)
98//! - [`MAX_FILE_SIZE`] — largest file accepted by the file loaders (std only)
99//!
100//! It also validates that the file is large enough to contain every tensor it claims,
101//! returning [`Error::ValidationError`] otherwise.
102//!
103//! ## Feature Flags
104//!
105//! - `std`: Enables file I/O ([`Reader::from_file`] / [`Writer::write_to_file`]) (default)
106
107extern crate alloc;
108
109mod base;
110mod reader;
111mod tensor;
112mod writer;
113
114#[cfg(feature = "std")]
115pub use base::MAX_FILE_SIZE;
116pub use base::{
117    Error, FORMAT_VERSION, HEADER_SIZE, Header, MAGIC_NUMBER, MAX_CBOR_RECURSION_DEPTH,
118    MAX_METADATA_SIZE, MAX_TENSOR_COUNT, MAX_TENSOR_SIZE, Scalar, ScalarConversionError,
119    TENSOR_ALIGNMENT, aligned_data_section_start,
120};
121pub use reader::Reader;
122pub use tensor::Tensor;
123pub use writer::Writer;
124
125/// The canonical file extension for burnpack files (without the leading dot).
126pub const EXTENSION: &str = "bpk";
127
128// Re-export the core types so callers can build [`Tensor`] entries and inspect descriptors
129// without depending on `burn-std` directly.
130pub use burn_std::{Bytes, DType, Shape};
burn_pack/lib.rs

burn_pack/
lib.rs
